读完这篇,你能做什么
能计算各 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.py
长行可在代码区域内左右滚动,按 Tab 可把焦点移入代码区。
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["各副本一致更新"]按顺序看,每一步做了什么
- 固定world size、rank编号及相同的初始参数。
- 检查每rank索引,明确尾部补齐或丢弃规则。
- 汇总当前更新有效目标数,推导反向损失系数。
- 与单进程大批次梯度做数值对照,再扩展真实进程组。
03对照实现代码
PyTorch v2.9.1 · DistributedSampler.__iter__;DistributedSampler.set_epoch
下例是shuffle=False路径的原创等价教学模型,不包含真实随机排列和分布式通信。
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只更新轮次,必须在创建这一轮迭代器前调用;它不自动保存完整训练游标或同步优化器。
代码先调用谁,再调用谁
- 训练循环在epoch开始调用
set_epoch __iter__生成共同索引并补齐或截断- 按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输出长度相同;交错还原后等于顺序全集的截断或循环前缀补齐。
只为当前rank生成索引,不构造全局补齐数组;每项常数时间计算。
保存为 Main.java,运行 javac --release 17 -encoding UTF-8 Main.java,再运行 java Main。程序包含固定输入自测,断言失败时会抛出错误。
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(训练进程)的样本索引,说明补齐与丢弃的代价,并推导有效目标数不同时的梯度缩放。
本篇重点加练
全局分母与集体通信顺序共同决定分布式目标,重点检查异步事件和空批次。
- 全局分母加权
- 集体通信与累积顺序
重点内容会多练几次。先说清楚怎么做,再试试空数据、出错或条件改变时会发生什么。
先自己做,卡住时记下问题,再回来查。问答题写清“为什么、什么时候适用”;动手题留下输入、输出和运行步骤,方便下次重做。
基础必做
- 作业 01代码推演
两份模型副本具有相同初始参数,但各处理不同样本。画出一次梯度同步到参数更新的事件图,并区分 rank、本地批次与全局批次。
- 提交内容
- 事件图与状态表
- 怎样算完成
- 副本与分片不混淆
- 同步对象明确
练到的知识点rank与world size · 模型副本和梯度同步 · 默认梯度平均
- 作业 02原理问答
数据含 7 个样本、world size=3。按明确声明的 sampler 设置手工追踪两轮索引,比较补齐、截断与 epoch 随机种子。
- 提交内容
- 采样轨迹表
- 怎样算完成
- 重复或丢弃可见
- 全局及本地顺序均有记录
练到的知识点DistributedSampler补齐与截断 · set_epoch与共同随机顺序
- 作业 03原理问答
三个 rank 中一个当前无有效标签,但其他 rank 仍有梯度。设计该步的参与契约,分别声明数值贡献与必须参与的通信事件。
- 提交内容
- 各 rank 时序及目标定义
- 怎样算完成
- 零贡献和不参与区分
- 通信序列一致性可检查
练到的知识点全局分母加权 · 集体通信与累积顺序
综合应用
- 作业 04编程实践
用 CPU 小模型模拟三个 rank,有效目标数分别为 2、0、5;比较聚合梯度与同一批数据单进程目标的梯度。
- 提交内容
- 模拟代码与差异记录
- 怎样算完成
- 全局分母可核对
- 模拟不声称验证通信实现
练到的知识点默认梯度平均 · 样本数与有效Token数 · 全局分母加权
- 作业 05排错实验
在 CPU 单机可用时运行两进程 Gloo 小实验,固定三个同步步骤并故意制造不一致分支;设置超时并保留各 rank 事件日志。
- 提交内容
- 脚本、环境及日志
- 怎样算完成
- 超时有界不会无限等待
- 故障注入不影响外部进程
练到的知识点set_epoch与共同随机顺序 · 集体通信与累积顺序
- 作业 06面试场景
分布式评估结果比单进程略高。设计按样本身份与有效目标计数审计的流程,覆盖尾部补齐和局部均值汇总。
- 提交内容
- 审计报告格式与样例
- 怎样算完成
- 重复计数可定位
- 训练与评估策略区分
练到的知识点DistributedSampler补齐与截断 · 样本数与有效Token数 · 全局分母加权
- 作业 07编程实践
用 CPU 日志模拟四步分布式累积,其中每步各 rank 有效计数不同;审计全局分母及通信事件匹配,注入一次漏同步验证检测。
- 提交内容
- 模拟检查器与故障日志
- 怎样算完成
- 数值账本与通信账本分开
- 不声称模拟验证真实后端
练到的知识点全局分母加权 · 集体通信与累积顺序
- 作业 08设计评审
任务同时含两个损失头,两个头的有效标签数量不同。定义跨 rank 的联合目标与权重,说明一个全局样本数是否足够表达它。
- 提交内容
- 双目标公式与统计字段
- 怎样算完成
- 每个损失头归约独立
- 全局权重含义明确
练到的知识点全局分母加权
进阶挑战 · 掌握前置后选做
- 作业 09设计评审
为梯度累积加同步抑制设计验证协议,先做 CPU 数值等价再选做真实多设备通信测量;明确哪些结论必须来自实际通信日志。
- 提交内容
- 两阶段验证计划
- 怎样算完成
- 数值与系统性能分开
- 设备前置条件明确
练到的知识点rank与world size · 默认梯度平均 · 集体通信与累积顺序
思维拓展
先写出你的猜想,再用例子或实验验证。答案可以不同,但要说清依据、适用情况,以及还有什么没弄懂。
- 作业 10思维拓展 · 条件变化
若某个 rank 永久退出,剩余进程继续训练意味着什么目标与样本分布变化?提出恢复或重启协议,并列出单机模拟无法验证的集体通信保证。
- 提交内容
- 弹性条件分析与待验证清单
- 怎样算完成
- 成员变化及样本归属明确
- 真实通信结论需要实际证据
练到的知识点全局分母加权 · 集体通信与累积顺序
完成后,再检验一次
给每题标记“独立完成 / 查阅后完成 / 尚未完成”。查过资料或没做出的题,记下卡在哪里。隔一天换一组输入再试,能解释结果和出错原因,再把这个知识点标为掌握。
继续阅读
先看能帮你完成当前练习的部分。例子跑通后,再读更深入的解释。
- 必读参考文档英文
PyTorch2.9 DDP 文档 (新标签页)↗
核对梯度平均、输入分片责任和
no_sync边界。 - 深入参考文档英文
PyTorch2.9.1 DistributedSampler (新标签页)↗
核对共同随机种子、尾部补齐与rank切片。
- 深入参考文档英文
PyTorch2.9.1 DDP 源码 (新标签页)↗
追踪梯度同步与参数副本更新约定。