PyTorch DDP 单卡改双卡训练结果对齐实战指南

单卡跑通的训练脚本直接套上torchrun --nproc_per_node2就一定能得到和单卡一致的结果吗我一开始也是这么以为的直到某次实验里 loss 曲线在双卡下明显抖了一下排查了大半天才发现是 DataLoader 的 shuffle 种子没对齐。LLM Training Lab 这个系列做到第 12 期我想专门把单卡改双卡 DDP 且结果不变这件事讲透——它看起来只是加几行DistributedDataParallel的包装实际上涉及随机种子、数据切分、梯度归约、BatchNorm 统计量、日志打印等一整条链路。这篇内容适合已经能用 PyTorch 跑通单卡训练、准备上多卡但不想让实验结果变味的读者也适合那些被 DDP 结果对不齐折磨过的同行。下面我按自己踩过的顺序把每个环节拆开讲。1. 先搞清楚结果不变到底指什么很多人一上来就问双卡和单卡 loss 能不能一模一样这个问题本身问得不够精确。DDP 的数学语义和单卡并不完全等价我们得先把不变拆成几个层次才知道哪些能对齐、哪些注定有差异。1.1 三个层次的对齐目标我把结果不变分成三档从易到难第一档训练能收敛最终指标在同一量级。这是最低要求绝大多数场景够用。第二档每个 step 的 loss 数值在浮点误差范围内一致。这要求数据顺序、梯度归约方式、随机数消耗完全对齐。第三档bit-wise 完全一致。这个基本做不到也不该追求因为 NCCL 的 all-reduce 归约顺序和单卡的累加顺序不同浮点加法不满足结合律末位必然有差异。我个人的经验是把目标定在第二档。也就是 loss 曲线肉眼重合、最终 checkpoint 的权重差异在 1e-5 量级以内。这个目标既有实际意义保证实验可复现又不会陷入无谓的 bit 级较劲。1.2 为什么 DDP 天然会引入差异要理解差异从哪来得先明白 DDP 干了什么。单卡时一个 batch 的梯度直接由这一份数据算出。双卡 DDP 时每张卡各拿一半数据算出各自的梯度然后通过 NCCL 做一次 all-reduce 求平均再各自更新。这里有两个关键点数据被切分了。原本一个 batch 的 32 条样本现在每卡 16 条。如果你的 batch size 设置没跟着调整等效 batch 就变了。梯度是先分后合的。单卡是32 条一起算梯度双卡是16 条算完求平均再和另一张卡的 16 条求平均。数学上如果 loss 是样本均值这两者等价但如果 loss 里有跨样本的项比如对比学习里的 in-batch negative切分后就完全不是一回事了。提示判断你的任务能不能无损切分就看 loss 是否可以写成每个样本独立贡献之和再取平均。能就能对齐不能就得改 loss 设计。1.3 一个容易忽略的前提等效 batch size假设单卡时batch_size32你改成双卡后如果每卡还是batch_size32那等效 batch 变成了 64。这时候 loss 曲线和单卡对不上是必然的因为优化轨迹都变了。正确做法是保持全局 batch size 不变单卡 32双卡就每卡 16。这样每个 step 看到的样本总数一致梯度期望一致学习率也不用动。我见过太多人在这里翻车然后去怀疑 DDP 本身有问题其实是 batch size 悄悄翻倍了。2. 随机种子对齐结果的第一道关卡数据顺序和初始化权重是训练里最大的两个随机源。单卡改双卡如果种子处理不当这两处都会错位loss 曲线自然对不上。2.1 全局种子的设置与陷阱PyTorch 里设置种子的标准写法是import torch import numpy as np import random def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)单卡时这样写没问题。但双卡时有个坑torch.cuda.manual_seed_all会给所有 GPU 设同一个种子这本身没错但如果你在每张卡上跑同样的初始化代码模型初始权重确实能一致。真正的问题出在DataLoader 的 worker 种子上。DataLoader 如果开了num_workers0每个 worker 会用自己的随机状态去 shuffle 和做数据增强。PyTorch 默认给每个 worker 分配的种子是基于base_seed worker_id而base_seed又和主进程的随机状态有关。双卡时两个进程的主进程种子如果不同worker 种子就不同数据顺序就乱了。2.2 让每张卡的数据切分可复现我的做法是给每个 rank 显式设置种子并且让 DataLoader 的 sampler 也带上 rank 信息import os import torch.distributed as dist def setup_ddp(): dist.init_process_group(backendnccl) rank dist.get_rank() local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) # 关键每个 rank 用不同但确定的种子 set_seed(42 rank) return rank, local_rank注意这里set_seed(42 rank)而不是所有 rank 都用 42。为什么因为如果用同一个种子两张卡会生成完全相同的随机数序列配合 DistributedSampler 的切分反而可能让某些数据增强重复。用42 rank保证每个 rank 的随机流独立且可复现。但这里又有个矛盾模型初始化我们希望所有 rank 一致DDP 会自动 broadcast rank 0 的参数所以其实初始化不一致也没关系会被覆盖。所以更稳妥的做法是模型初始化用固定种子 42数据相关的随机用 42rank。分开管理逻辑更清晰。2.3 DistributedSampler 的 shuffle 行为DistributedSampler是 DDP 数据切分的核心。它的工作方式是每个 epoch 用seed epoch作为随机种子打乱整个数据集然后按 rank 轮流分配样本。from torch.utils.data.distributed import DistributedSampler sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue, seed42, # 固定种子 drop_lastTrue # 建议开启避免各卡样本数不均 )这里seed42是固定的配合set_epoch(epoch)在每个 epoch 开始时调用就能保证每个 epoch 的切分确定。drop_lastTrue很重要——如果数据集大小不能被 world_size 整除最后几个样本会导致各卡数量不一致DDP 的 all-reduce 会卡住或报错。注意DistributedSampler的 shuffle 和单卡RandomSampler的 shuffle 结果是不一样的。单卡是整个数据集打乱后顺序取双卡是打乱后按 rank 轮流取。所以即使种子相同双卡每个 step 看到的样本组合也和单卡不同。这是第二档对齐里最难处理的一环后面第 4 节会专门讲怎么绕。3. 模型包装与梯度归约的细节数据对齐了接下来是模型侧。DDP 的包装看似一行代码但里面的 bucket、find_unused_parameters 等参数会直接影响结果和性能。3.1 DDP 包装的正确姿势标准写法from torch.nn.parallel import DistributedDataParallel as DDP model MyModel().to(local_rank) model DDP( model, device_ids[local_rank], output_devicelocal_rank, find_unused_parametersFalse, # 默认 False别乱开 broadcast_buffersTrue, # 默认 True gradient_as_bucket_viewTrue # 省显存 )几个参数逐个说find_unused_parameters如果你的模型有分支某些参数在部分 step 不参与计算设 True 能避免报错但会拖慢速度每个 step 都要遍历计算图找未使用参数。LLM 训练里一般所有参数都参与保持 False。broadcast_buffers控制是否在 forward 前把 rank 0 的 buffer如 BatchNorm 的 running_mean广播到其他卡。LLM 基本不用 BatchNorm但如果你有自定义 buffer这个参数很关键。gradient_as_bucket_view让梯度直接写进 DDP 的通信 bucket省一份显存拷贝。开了之后不能再对梯度做原地修改注意兼容性。3.2 梯度归约的数学等价性DDP 默认对梯度做平均sum 后除以 world_size。这对应的是 loss 取 mean 的情况。如果你的 loss 是 sum 而不是 mean那 DDP 的平均就和单卡的 sum 不一致了等效学习率会变成原来的 1/world_size。解决办法有两个把 loss 改成 mean推荐最省事。在 loss 上乘以 world_size 补偿。我一般选第一种。因为 mean 的 loss 数值不随卡数变化日志看起来也直观。如果你非要用 sum记得在loss.backward()前乘world_size否则双卡的学习率等效减半收敛会明显变慢。3.3 梯度累积与 DDP 的配合LLM 训练常用梯度累积来模拟大 batch。单卡时是累积 N 个 step 再更新双卡时如果还按原来的 N等效 batch 又翻倍了。正确做法梯度累积步数除以 world_size。单卡累积 8 步双卡就累积 4 步保证全局等效 batch 不变。同时注意 loss 要除以累积步数这个除法在单卡双卡下逻辑一致不用改。accum_steps 8 // world_size for i, batch in enumerate(loader): loss model(batch) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()4. 让双卡每个 step 的样本组合逼近单卡这是第二档对齐里最硬的一块。前面说过DistributedSampler 的切分方式和单卡 RandomSampler 不同导致每个 step 的样本组合对不上。如果你追求 loss 曲线逐 step 重合就得想办法让双卡的每卡 batch 拼接等于单卡的 batch。4.1 用自定义 Sampler 模拟单卡顺序思路是先用单卡的 RandomSampler 逻辑生成一个全局的样本索引序列然后按 step 切分每个 step 的前半给 rank 0后半给 rank 1。class SingleCardOrderSampler(torch.utils.data.Sampler): def __init__(self, dataset_size, batch_size, world_size, rank, seed42, epoch0): self.dataset_size dataset_size self.batch_size batch_size self.world_size world_size self.rank rank self.seed seed self.epoch epoch def __iter__(self): g torch.Generator() g.manual_seed(self.seed self.epoch) # 复现单卡的全局打乱 indices torch.randperm(self.dataset_size, generatorg).tolist() # 按 step 切分每个 step 取 batch_size 个 per_rank self.batch_size // self.world_size result [] for start in range(0, len(indices) - self.batch_size 1, self.batch_size): step_indices indices[start:start self.batch_size] # 当前 rank 取自己那一段 offset self.rank * per_rank result.extend(step_indices[offset:offset per_rank]) return iter(result) def __len__(self): return self.dataset_size // self.batch_size * (self.batch_size // self.world_size)这样每个 step 里rank 0 拿的是单卡 batch 的前半rank 1 拿的是后半all-reduce 平均后和单卡整 batch 的梯度在数学上等价浮点误差除外。实测下来 loss 曲线能贴得很近。4.2 数据增强的随机性对齐如果数据集有随机增强比如图像裁剪、文本 mask每个样本的增强结果也依赖随机状态。上面的 sampler 只对齐了取哪些样本没对齐怎么增强。要完全对齐得让每个样本的增强种子只和样本 id 有关而不是和 worker 的随机流有关。做法是在__getitem__里用样本 id 派生种子def __getitem__(self, idx): rng random.Random(self.base_seed idx) # 用 rng 做增强保证同一样本在任何 rank、任何 worker 下增强结果一致 ...这样无论样本被分到哪张卡增强结果都一样双卡和单卡的差异就只剩浮点归约误差了。4.3 什么时候不必强求逐 step 对齐说实话上面这套自定义 sampler 有维护成本。如果你的实验只是看最终指标不要求 loss 曲线逐 step 重合那用标准 DistributedSampler 就够了。我自己的判断标准是调参阶段用标准 sampler看趋势就行。复现论文/对比实验用自定义 sampler保证严格对齐。生产训练用标准 sampler性能优先。5. 那些不影响 loss 但会让你怀疑人生的坑有些问题不会改变数学结果但会让你的日志、checkpoint、评估乱套误以为 DDP 出错了。5.1 日志重复打印DDP 下每个 rank 都会执行你的 print 语句于是同一行日志打印了 world_size 次。解决办法是只在 rank 0 打印if rank 0: print(fstep {i}, loss {loss.item():.4f})但注意loss.item()会触发一次 GPU 同步频繁调用拖慢训练。更好的做法是累积几个 step 再打印或者用loss.detach()存下来最后统一处理。5.2 checkpoint 保存的竞争如果每个 rank 都执行torch.save会互相覆盖甚至写坏文件。标准做法是只在 rank 0 保存并且保存前确保所有 rank 的模型状态一致DDP 保证参数一致但 buffer 如果没 broadcast 可能不一致。if rank 0: torch.save({ model: model.module.state_dict(), # 注意用 .module 去掉 DDP 包装 optimizer: optimizer.state_dict(), epoch: epoch, }, ckpt.pt) dist.barrier() # 其他 rank 等 rank 0 存完用model.module.state_dict()而不是model.state_dict()否则 key 里会多出module.前缀加载时对不上。5.3 评估指标的跨卡聚合验证时每张卡算自己那部分数据的指标最后要聚合。简单平均是错的因为各卡样本数可能不同。正确做法是按样本数加权或者用dist.all_reduce把分子分母分别求和correct torch.tensor(correct_count, devicelocal_rank) total torch.tensor(total_count, devicelocal_rank) dist.all_reduce(correct, opdist.ReduceOp.SUM) dist.all_reduce(total, opdist.ReduceOp.SUM) acc correct.item() / total.item()这个坑我在早期项目里踩过双卡评估准确率比单卡低了两个点查了半天才发现是简单平均导致的。6. 启动脚本与 NCCL 环境最后聊聊启动方式和通信后端。这部分配错了训练要么起不来要么慢得离谱。6.1 torchrun 的参数含义现在推荐用torchrun而不是老的python -m torch.distributed.launchtorchrun \ --nproc_per_node2 \ --nnodes1 \ --node_rank0 \ --master_addr127.0.0.1 \ --master_port29500 \ train.pynproc_per_node每台机器的进程数通常等于 GPU 数。master_addr/portrank 0 的通信地址单机多卡用本地回环即可。脚本里通过os.environ[LOCAL_RANK]拿当前进程对应的 GPU。6.2 NCCL 的几个关键环境变量NCCL 是 NVIDIA 显卡的默认通信后端几个变量值得关注变量作用建议值NCCL_DEBUG调试日志级别排查时设 INFO平时不设NCCL_IB_DISABLE禁用 InfiniBand单机无 IB 时设 1NCCL_P2P_DISABLE禁用 GPU 间 P2P遇到 P2P 报错时设 1NCCL_SOCKET_IFNAME指定网卡多网卡机器需指定单机双卡一般不用动这些NCCL 会自动选最优路径。但如果你的机器有多张网卡NCCL_SOCKET_IFNAME不指定可能导致通信走错网卡速度骤降。我遇到过一台机器双卡训练比单卡还慢最后发现是 NCCL 走了管理网卡。6.3 验证 DDP 是否真的生效跑起来后怎么确认 DDP 在工作看两个地方启动日志里应该有Initializing process group with backend nccl和world_size2。用nvidia-smi看两张卡的显存占用和利用率应该都在动。如果只有一张卡在跑多半是local_rank没设对或者模型没.to(local_rank)。7. 一套可复现的对照实验流程讲了这么多原理最后给一套我自己常用的验证流程帮你确认双卡改造是否成功。7.1 小数据快速对照别一上来就跑全量。先取 100 个样本单卡跑 50 步记录每步 loss再用双卡跑 50 步对比两条曲线。如果前 10 步就明显分叉说明种子或数据切分有问题如果只是末位抖动那是正常的浮点误差。7.2 逐项排查清单对照下面这张表逐项检查基本能覆盖 90% 的对不齐问题检查项单卡设置双卡应改为全局 batch size32每卡 16全局仍 32梯度累积步数84除以 world_size随机种子42模型 42数据 42rankloss 归约meanmeanDDP 自动平均checkpoint直接存仅 rank 0 存用 .module日志直接 print仅 rank 0 print7.3 我踩过的最隐蔽的坑有一次双卡 loss 比单卡高了一截排查了两小时最后发现是DataLoader的num_workers在双卡下设成了 4而单卡是 2。worker 数量变了数据预取的随机状态消耗顺序也变了导致增强结果错位。这个坑的教训是对照实验时除了必须改的参数其他一律保持一致包括 num_workers、pin_memory、persistent_workers 这些看似无关的选项。还有一次是优化器的问题。单卡用的是AdamW双卡时我手滑写成了Adamweight decay 行为不同loss 曲线自然对不上。这种低级错误在改代码时特别容易发生建议改完后用diff对比一下单卡和双卡的配置文件。把上面这些环节都对齐之后双卡和单卡的 loss 曲线基本能重合到小数点后三位最终 checkpoint 的权重差异在 1e-5 量级。这个精度对绝大多数实验来说已经足够剩下的差异就是 NCCL 归约顺序带来的浮点末位误差属于物理极限不用再纠结。真正要花心思的是数据切分和随机种子这两块它们才是决定结果变不变的关键。