qinghuan
← 预训练与训练工程 / 全部笔记

单元 28 / 预训练与训练工程

P1 · 重点提升L2 · 进阶

DDP、梯度 AllReduce 与分布式采样

先看多张设备各自处理哪些样本、最后怎样合并梯度,再用采样索引和梯度对照理解 DDP 副本同步,以及有效目标数不同的权重。

本篇目录 · 按需要跳转

读完这篇,你能做什么

能计算各 rank(训练进程)的样本索引,说明补齐与丢弃的代价,并推导有效目标数不同时的梯度缩放。

学习理由、准备要求与知识点

为什么学它

多设备都跑起来了,仍可能因为重复采样、平均时分母算错或通信顺序不一致而出问题。

先读这些内容

知识点清单

  • rank与world size
  • 模型副本和梯度同步
  • 默认梯度平均
  • DistributedSampler补齐与截断
  • set_epoch与共同随机顺序
  • 样本数与有效Token数
  • 全局分母加权
  • 集体通信与累积顺序

运行前先确认Java17、PyTorch2.9.1 CPU。实验真实运行DistributedSampler与autograd,但只在一个进程手工模拟梯度平均;未初始化进程组、执行网络AllReduce或多机训练。

01先弄懂它是什么

有多块设备时,可以让它们各算一部分数据,再汇总梯度,这就是 DDP 数据并行。每个工作进程都持有一份模型,rank 是它在通信组里的编号,world size 是进程总数 W。各 rank 分的是数据,不是不同网络层;若都读同一批数据,就可能只是在重复工作。模型初始状态和更新规则也要保持一致,才能在同步梯度后继续得到相同参数。

汇总梯度时,要确认每个进程的平均数代表多少数据。默认 DDP 的处理相当于跨 rank 平均梯度,自定义通信钩子等设置可能改变这点。若各 rank 都对同样数量的有效目标取平均,再平均梯度就对应全局均值;若目标数量不同,直接平均各自均值,会让目标少的那部分占过大权重。语言模型还要数未被屏蔽的词元,不能只看输入有几行。

分布式采样器做的是分发数据编号,并不执行模型计算。假设七条数据分给三个 rank,数量分不齐:可以重复前面的两条,凑成九条,每个 rank 拿三条;也可以丢掉末尾一条,每个拿两条。前者重复使用数据,后者本轮漏掉数据,验证和测试算总分时都要考虑。

这里是 DistributedSampler 的 drop_last 选择;DataLoader 的 drop_last 则决定每个 rank 最后凑不满一批时怎么办,它们是两层不同的截断。

文中这几个词,先认一认

AllReduce
所有参与进程提交形状相同的数据,按求和等规则合并,再让各进程拿到合并结果。梯度同步还可能包含平均缩放。
全局有效数
这一次更新中,所有 rank、所有微批里真正参与损失计算的位置总数。
采样补齐
重复一部分已有数据索引,把总长度补到进程总数的整数倍,使每个 rank 拿到同样多的索引。

还有陌生词?到术语小抄里看一个具体例子 ↗

02看看它怎么工作

合并梯度时,每个有效目标应有相同的份量,而不是每个 rank 各算一票。例如两个 rank 分别有一个和三个有效目标,第二个局部平均数应占四分之三。一般设 rank r 的目标数是 nr、损失总和是 Sr,所有 rank 合计 Ntotal 个目标,我们要优化的是 sum(Sr)/Ntotal。

若 DDP 会把 W 个梯度再平均一次,每个 rank 反向前就用 W×Sr/Ntotal;同步时除以 W,正好抵消。若已有局部均值 Lr,则先乘 W×nr/Ntotal。前提是计数准确、同步确实采用平均,之后不能再额外除一次 W。

各进程要从同一份打乱后的总顺序中分取数据,才能明确知道谁拿哪一项。开启 shuffle 时,每轮先调用 set_epoch,让种子和轮次一起决定新顺序;若每个 rank 各自随意重排全集,就不能保证它们拿到的数据互不重复。Java 实验没有实现打乱,只展示顺序输入如何补齐、截断和跨步切片,因此不声称复现了 PyTorch 的随机数发生器。

Python实验构造两个虚拟rank,一个样本和三个样本。它分别对相同标量参数求梯度,对照局部均值简单平均、按目标数加权以及全量参考。后两者一致,前者不一致。随后调用真实DistributedSampler核对七个样本三rank时补齐到九项。这个实验验证数学与采样API,不涉及网络性能、进程故障或通信死锁恢复。

先跑通一个最小例子

运行 experiment.py。两组梯度在同一CPU进程计算,DistributedSampler显式传num_replicas和rank,因此无需进程组。

展开收起代码 · experiment.pypython · 35 行

长行可在代码区域内左右滚动,按 Tab 可把焦点移入代码区。

experiment.py · python
import torch
from torch.utils.data import DistributedSampler
torch.set_num_threads(1)
torch.manual_seed(28)
x = torch.tensor([1., 2., 3., 4.], dtype=torch.float64)
def grad(values, scale=1.0, reduction="mean"):
    w = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
    errors = (w * values).square()
    loss = errors.mean() if reduction == "mean" else errors.sum()
    return torch.autograd.grad(loss * scale, w)[0]
g0, g1 = grad(x[:1]), grad(x[1:])
reference = grad(x)
# 不同有效数的局部均值,不能无权重平均
wrong = (g0 + g1) / 2
weighted = (g0 * 1 + g1 * 3) / 4
assert not torch.allclose(wrong, reference)
assert torch.allclose(weighted, reference)
# 模拟DDP平均:每rank损失和先乘W/Ntotal
scaled = (grad(x[:1], 2/4, "sum") + grad(x[1:], 2/4, "sum")) / 2
assert torch.allclose(scaled, reference)
def indices(n, world, drop):
    return [list(DistributedSampler(list(range(n)), num_replicas=world,
                 rank=r, shuffle=False, drop_last=drop)) for r in range(world)]
padded = indices(7, 3, False)
assert padded == [[0,3,6], [1,4,0], [2,5,1]]
assert sum(map(len, padded)) == 9
assert indices(7, 3, True) == [[0,3], [1,4], [2,5]]
assert indices(2, 4, False) == [[0], [1], [0], [1]]
# 所有rank同seed和epoch可重建相同全局排列
a = DistributedSampler(list(range(12)), num_replicas=2, rank=0, seed=9)
a.set_epoch(3)
b = DistributedSampler(list(range(12)), num_replicas=2, rank=0, seed=9)
b.set_epoch(3)
assert list(a) == list(b)
print("8 sampler and weighted gradient checks passed")

核对运行结果真实采样索引和手工加权梯度均通过8项检查;不报告网络AllReduce已运行。

正在加载流程图…

损失分母必须与梯度同步方式匹配;图示省略实际通信协议。
查看流程图源码(Mermaid)
flowchart TD
 A["共同模型与采样顺序"] --> B["rank0 本地目标"]
 A --> C["rank1 本地目标"]
 B --> D["本地损失和与计数"]
 C --> E["本地损失和与计数"]
 D --> F["全局有效数与正确缩放"]
 E --> F
 F --> G["梯度同步平均"]
 G --> H["各副本一致更新"]

按顺序看,每一步做了什么

  1. 固定world size、rank编号及相同的初始参数。
  2. 检查每rank索引,明确尾部补齐或丢弃规则。
  3. 汇总当前更新有效目标数,推导反向损失系数。
  4. 与单进程大批次梯度做数值对照,再扩展真实进程组。

03对照实现代码

DistributedSampler 的补齐与按rank切片 (新标签页)↗

PyTorch v2.9.1 · DistributedSampler.__iter__;DistributedSampler.set_epoch

下例是shuffle=False路径的原创等价教学模型,不包含真实随机排列和分布式通信。

核心机制代码 · python
def rank_indices(n, world, rank, drop=False):
    if n < 0 or world <= 0 or not 0 <= rank < world:
        raise ValueError("invalid sampler")
    # 补齐用上取整,丢弃用下取整
    per_rank = n // world if drop else (n + world - 1) // world
    total = per_rank * world
    if total == 0:
        return []
    # 前缀循环补齐;丢弃时total不超过n,所以取模不改变索引
    global_indices = [i % n for i in range(total)]
    # 同一全局排列按rank跨步分配
    return global_indices[rank:total:world]

真实实现以seed+epoch产生共同randperm,而不是每个rank独立重排自己的候选全集。

drop_last=False按需要重复前缀;数据比rank数还少时,也支持多次重复后再裁到所需长度。

完成总长度调整后使用rank起点和num_replicas步长切片,保证每个rank的样本数一致。

set_epoch只更新轮次,必须在创建这一轮迭代器前调用;它不自动保存完整训练游标或同步优化器。

代码先调用谁,再调用谁

  1. 训练循环在epoch开始调用set_epoch
  2. __iter__生成共同索引并补齐或截断
  3. 按rank跨步切片后交给DataLoader

04动手写一个例子

不打乱时的分布式索引规划

输入样本总数n、进程数w、rank和是否丢弃尾部,返回该rank索引。省略随机打乱但保留补齐与截断语义。

输入 1
n=7,w=3,rank=1,drop=false
输出 1
[1,4,0]

先计算每rank应有的元素个数。补齐选择上取整,丢弃选择下取整,统一决定各rank参与的数据量。

该rank第j个索引对应全局位置rank+j*w;补齐阶段全局位置对n取模即可复现顺序前缀循环。

使用long计算全局位置,避免虽然返回数组不大但rank与步长乘积溢出int。空数据直接返回空数组,绕开对零取模。

为什么这样写是对的所有rank输出长度相同;交错还原后等于顺序全集的截断或循环前缀补齐。

时间 O(ceil(n/w))额外空间 O(ceil(n/w))输出,O(1)额外空间

只为当前rank生成索引,不构造全局补齐数组;每项常数时间计算。

保存为 Main.java,运行 javac --release 17 -encoding UTF-8 Main.java,再运行 java Main。程序包含固定输入自测,断言失败时会抛出错误。

Main.java · 完整可运行程序 · Java
import java.util.*;
public class Main {
    static int[] sample(int n,int world,int rank,boolean drop){
        if(n<0 || world<=0 || rank<0 || rank>=world)
            throw new IllegalArgumentException();
        // 用long上取整避免n+world-1溢出
        long count=drop ? n/world : ((long)n+world-1)/world;
        int[] result=new int[(int)count];
        if(n==0)return result;
        for(int j=0;j<result.length;j++){
            // 先算全局位置,补齐重复顺序前缀
            long position=rank+(long)j*world;
            result[j]=(int)(position%n);
        }
        return result;
    }
    static void check(boolean ok){if(!ok)throw new AssertionError();}
    public static void main(String[] args){
        check(Arrays.equals(sample(7,3,1,false),new int[]{1,4,0}));
        check(Arrays.equals(sample(7,3,2,true),new int[]{2,5}));
        check(Arrays.equals(sample(2,4,3,false),new int[]{1}));
        check(sample(2,4,0,true).length==0);
        check(sample(0,3,2,false).length==0);
        // 大样本总数但world也大,本rank仅产生两个索引
        check(Arrays.equals(sample(Integer.MAX_VALUE,Integer.MAX_VALUE-1,0,false),
              new int[]{0,Integer.MAX_VALUE-1}));
        try{sample(3,2,2,false);throw new AssertionError();}
        catch(IllegalArgumentException expected){}
        System.out.println("7 rank-index checks passed");
    }
}

这些特殊情况也要试一试

  • world大于样本数时补齐会重复,丢弃可能让所有rank为空。
  • n等于零不能执行取模。
  • world为零或rank越界应在通信前失败。
进阶选读 · 不补齐评测与去重聚合

场景评测要求每个样本只贡献一次指标,但简单DistributedSampler补齐造成重复计数。

做法使用不重复的分片,按稳定样本ID聚合结果;若各rank步数不同,避免在不一致的循环中调用需要全体参与的集体操作,在统一阶段汇总指标。

复杂度索引分片O(n),去重聚合需要O(n)标识或外部存储;通信量由预测结果大小决定。

好处与代价评测准确计数和训练等步同步是不同约束,不能盲目复用同一循环。

05项目里怎么用

DDP不会替应用自动切分输入。启动之后应先记录每rank少量样本ID,核对重叠和尾部行为,再比对一个更新步的参数变化。固定参数、关闭随机层、使用小批次做数值对照,比直接观察长期损失曲线更适合发现错误缩放。之后再逐一恢复dropout、混合精度和梯度累积。

做梯度累积时,可以在中间几个微批使用 no_sync 暂缓同步,最后一次再正常同步;前向和反向都要放进对应上下文,不能只包住 backward。计算有效目标数时,也要覆盖整个累积窗口。有的 rank 即使一个有效目标也没有,仍需按约定参加必要的集体通信;若它单独跳过,其他进程可能一直等不到它。

要事先分清各 rank 谁写日志、谁保存检查点、谁汇总评测。只让一个 rank 写文件能减少重复写入,但不表示其他 rank 的恢复信息可以丢掉:它们的随机数状态、数据读到哪里,以及恢复所用的配置,仍要保存或能重新构造。本文没有运行多进程通信,真正部署时还要单独核验通信后端、初始化超时、各进程参加通信的顺序,以及一个节点退出后怎样让整项训练正确失败。

容易踩的坑

  • 所有rank读取同一数据且没有分片。
  • 忘记set_epoch,使各轮重复同一采样顺序。
  • 不同Token数局部均值直接平均,改变目标权重。
  • 某个rank提前跳过集体操作,导致其他rank等待。

几种做法,怎么选

表格有四列,窄屏可左右滑动查看代价和适用条件。

方案优势代价与限制适用条件
补齐采样训练步数容易一致尾部样本重复明确接受重复权重的训练
丢弃尾部没有补齐重复本轮部分样本不被处理样本多且轮次打乱的训练
不等长分片评测样本只算一次需要处理不同步数与聚合严格计数的评测

06检验是否读懂了

先自己想、自己写,再展开答案核对。场景题是假设的练习情境;遇到还没学过的内容,可以做完基础后再回来。

先用自己的话解释

01DDP是不是把模型层分到不同rank?

展开参考答案与推导

普通DDP每个rank持有模型副本,分配的是数据;参数分片或流水线并行是其他机制。

02本地mean再跨rank平均何时正确?

展开参考答案与推导

有效目标数相同且同步采用平均、无其他改变权重的操作时,可对应全局目标mean。

03为什么补齐会影响评测?

展开参考答案与推导

重复索引会重复贡献指标,不去重时某些样本权重更高。

换个情况,还会处理吗

01两个rank分别有1和3个有效目标怎么办?

展开参考答案与推导

汇总全局有效数4,各rank损失和乘2/4后反向;DDP再平均两者,得到全局每目标均值梯度。

02rank0没有效标签能直接continue吗?

展开参考答案与推导

不能独立改变同步顺序;应设计统一的空目标处理并参与必要同步,避免除零及进程等待。

03增加GPU后效果变差但代码无报错怎么查?

展开参考答案与推导

核对全局批次、学习率日程、样本补齐、有效Token缩放,再与单进程参考更新比较。

写代码试一试

01算补齐索引。

输入 / 条件
n=5,w=2,rank=1,shuffle=false
输出 / 预期
[1,3,0]

提示 先补到6项再隔位取。

展开参考答案与推导

全局序列为0,1,2,3,4,0,rank1取位置1,3,5。

02合并两个局部梯度。

输入 / 条件
g0=2,n0=1;g1=6,n1=3
输出 / 预期
5

提示 按目标数加权。

展开参考答案与推导

(2×1+6×3)/4=5;直接平均是4。

03实现对空数据安全的采样。

输入 / 条件
n=0,w=4,rank=2
输出 / 预期
[]

提示 不要先进行取模。

展开参考答案与推导

每rank长度为零,立即返回空数组。

读完之后,独立完成

课后作业:自己做一遍

本篇共 10 道作业,先完成 8 道基础与综合任务,再按前置知识挑战进阶题。其中包含 1 道思维拓展,试着换个条件,看看原来的办法是否还管用。不附参考答案。

这次练什么

能计算各 rank(训练进程)的样本索引,说明补齐与丢弃的代价,并推导有效目标数不同时的梯度缩放。

本篇重点加练

全局分母与集体通信顺序共同决定分布式目标,重点检查异步事件和空批次。

  • 全局分母加权
  • 集体通信与累积顺序

重点内容会多练几次。先说清楚怎么做,再试试空数据、出错或条件改变时会发生什么。

先自己做,卡住时记下问题,再回来查。问答题写清“为什么、什么时候适用”;动手题留下输入、输出和运行步骤,方便下次重做。

基础必做

  1. 作业 01代码推演

    两份模型副本具有相同初始参数,但各处理不同样本。画出一次梯度同步到参数更新的事件图,并区分 rank、本地批次与全局批次。

    提交内容
    事件图与状态表
    怎样算完成
    • 副本与分片不混淆
    • 同步对象明确

    练到的知识点rank与world size · 模型副本和梯度同步 · 默认梯度平均

  2. 作业 02原理问答

    数据含 7 个样本、world size=3。按明确声明的 sampler 设置手工追踪两轮索引,比较补齐、截断与 epoch 随机种子。

    提交内容
    采样轨迹表
    怎样算完成
    • 重复或丢弃可见
    • 全局及本地顺序均有记录

    练到的知识点DistributedSampler补齐与截断 · set_epoch与共同随机顺序

  3. 作业 03原理问答

    三个 rank 中一个当前无有效标签,但其他 rank 仍有梯度。设计该步的参与契约,分别声明数值贡献与必须参与的通信事件。

    提交内容
    各 rank 时序及目标定义
    怎样算完成
    • 零贡献和不参与区分
    • 通信序列一致性可检查

    练到的知识点全局分母加权 · 集体通信与累积顺序

综合应用

  1. 作业 04编程实践

    用 CPU 小模型模拟三个 rank,有效目标数分别为 2、0、5;比较聚合梯度与同一批数据单进程目标的梯度。

    提交内容
    模拟代码与差异记录
    怎样算完成
    • 全局分母可核对
    • 模拟不声称验证通信实现

    练到的知识点默认梯度平均 · 样本数与有效Token数 · 全局分母加权

  2. 作业 05排错实验

    在 CPU 单机可用时运行两进程 Gloo 小实验,固定三个同步步骤并故意制造不一致分支;设置超时并保留各 rank 事件日志。

    提交内容
    脚本、环境及日志
    怎样算完成
    • 超时有界不会无限等待
    • 故障注入不影响外部进程

    练到的知识点set_epoch与共同随机顺序 · 集体通信与累积顺序

  3. 作业 06面试场景

    分布式评估结果比单进程略高。设计按样本身份与有效目标计数审计的流程,覆盖尾部补齐和局部均值汇总。

    提交内容
    审计报告格式与样例
    怎样算完成
    • 重复计数可定位
    • 训练与评估策略区分

    练到的知识点DistributedSampler补齐与截断 · 样本数与有效Token数 · 全局分母加权

  4. 作业 07编程实践

    用 CPU 日志模拟四步分布式累积,其中每步各 rank 有效计数不同;审计全局分母及通信事件匹配,注入一次漏同步验证检测。

    提交内容
    模拟检查器与故障日志
    怎样算完成
    • 数值账本与通信账本分开
    • 不声称模拟验证真实后端

    练到的知识点全局分母加权 · 集体通信与累积顺序

  5. 作业 08设计评审

    任务同时含两个损失头,两个头的有效标签数量不同。定义跨 rank 的联合目标与权重,说明一个全局样本数是否足够表达它。

    提交内容
    双目标公式与统计字段
    怎样算完成
    • 每个损失头归约独立
    • 全局权重含义明确

    练到的知识点全局分母加权

进阶挑战 · 掌握前置后选做

  1. 作业 09设计评审

    为梯度累积加同步抑制设计验证协议,先做 CPU 数值等价再选做真实多设备通信测量;明确哪些结论必须来自实际通信日志。

    提交内容
    两阶段验证计划
    怎样算完成
    • 数值与系统性能分开
    • 设备前置条件明确

    练到的知识点rank与world size · 默认梯度平均 · 集体通信与累积顺序

思维拓展

先写出你的猜想,再用例子或实验验证。答案可以不同,但要说清依据、适用情况,以及还有什么没弄懂。

  1. 作业 10思维拓展 · 条件变化

    若某个 rank 永久退出,剩余进程继续训练意味着什么目标与样本分布变化?提出恢复或重启协议,并列出单机模拟无法验证的集体通信保证。

    提交内容
    弹性条件分析与待验证清单
    怎样算完成
    • 成员变化及样本归属明确
    • 真实通信结论需要实际证据

    练到的知识点全局分母加权 · 集体通信与累积顺序

完成后,再检验一次

给每题标记“独立完成 / 查阅后完成 / 尚未完成”。查过资料或没做出的题,记下卡在哪里。隔一天换一组输入再试,能解释结果和出错原因,再把这个知识点标为掌握。

继续阅读

先看能帮你完成当前练习的部分。例子跑通后,再读更深入的解释。

打开资料库,找下一份学习材料 ↗