1. 为什么单卡不够用了并行训练的动机拆解我最早接触并行训练是因为一次很尴尬的经历。手上有个视觉任务模型参数量大概1.2亿输入分辨率256×256batch size想开到64结果单张24G显存的卡连batch size 16都撑不住直接OOM。当时第一反应是缩模型或者降batch但降batch会直接影响BN层的统计稳定性和最终的mAP缩模型又意味着精度掉点。换句话说不是模型不行是硬件边界到了。这就是并行训练存在的意义当单卡的显存和算力无法支撑目标配置时把计算和数据切分到多张卡上用通信换空间和速度。但这里有个关键认知很多人一开始会混淆——并行训练解决的是两个不同维度的问题显存维度单卡放不下模型或批次需要把参数、梯度、优化器状态拆分到多卡吞吐维度单卡跑得慢想让更多卡同时算缩短单个epoch的墙钟时间PyTorch在工程上给出的两条主要路径就是DPDataParallel和DDPDistributedDataParallel。前者是早期单进程多线程方案后者是现在的默认推荐方案。理解这两者的差异本质上是理解PyTorch的多卡通信模型以及为什么DDP能把DP拍在沙滩上。如果你只用过nn.DataParallel(model).cuda()这种一行改法那大概率踩过以下坑主卡显存莫名比其他卡高一大截、显存利用率上不去、多机根本没法用、训练后期速度反而变慢。这些问题都不是玄学背后全是通信拓扑和GIL的锅。我打算把这两条路线从头拆一遍包括原理、代码、参数选择、踩坑排查尽量把这套东西讲透让刚入门的人能直接抄作业让用过的人能补齐底层认知。1.1 先搞清楚数据并行到底并行了什么数据并行的核心思想非常朴素每张卡持有一份完整的模型副本喂给每张卡不同的数据子集各自算梯度然后把梯度汇总求平均再用平均梯度更新所有卡的参数。这样从数学上看等价于用了一个更大的batch size在大模型上做单卡训练前提是BN的处理要合理。用一个生活类比。假设有一个小组做同一道数学题每个人每张卡做不同的题型不同的数据分片做完后大家把解题心得梯度汇总到组长那里组长求平均后把统一答案更新的参数发回给每个人。每个人手上始终有完整的解题模板完整模型只是每次处理的题目片段不同。这个模式决定了它的两个天然特性一是显存占用不会因为卡变多而降低每张卡都有完整模型所以数据并行解决算得快不直接解决模型太大装不下二是通信量正比于参数量每步都要同步梯度所以参数量越大通信开销越重什么时候该上更高级的并行策略这个点很关键。1.2 单机多卡之前先确认你的硬件能做点什么在聊代码之前有几个前置条件必须打牢否则后面全是玄学问题检查项说明常见坑卡数一致性DP/DDP都要求所有参与训练的卡型号尽量一致混插不同显存/算力的卡DDP会以最慢的卡为瓶颈DP会以最小显存为上限通信带宽同一台机器内多卡走NVLink或PCIe跨机走网卡PCIe 4.0 x16理论带宽约32GB/s万兆网卡只有1.25GB/s跨机通信是瓶颈拓扑结构nvidia-smi topo -m能看卡间连接拓扑不均衡会导致DDP里某些rank通信慢拖累整体CUDA可见性CUDA_VISIBLE_DEVICES控制进程能看到的卡设置了环境变量又用device_ids乱指容易越界有一点经验值得单独拎出来说同机器内做数据并行优先选DDP而不是DP即使你只有一张卡也建议写DDP因为后续换机器、加卡、改多机时迁移成本几乎为零。DP的代码看着简洁但它挖的坑通常在训练跑起来之后才爆出来。2. DP的原理拆解一个进程管所有卡问题都藏在这里DP全称torch.nn.DataParallel它的设计非常省事——你原来的单卡代码几乎不用改套一层nn.DataParallel(model)就能跑。这种低门槛让它在早期传播很广但它的底层实现决定了几个绕不开的局限。2.1 单进程多线程GIL和主卡的宿命DP的工作机制可以拆成四步把模型复制到所有可见的GPU上每张卡一份在主进程里把输入batch按卡数切分分发给各卡各卡前向、反向计算得到各自梯度把各卡的梯度汇聚到device 0主卡求和取平均再更新主卡参数然后把更新后的参数广播回其他卡注意第三步和第四步问题就出在这所有卡的计算是由主进程通过Python多线程调度的。Python有GIL全局解释器锁同一时刻一个进程只有一个线程在执行Python字节码。虽然GPU计算本身发生在C/CUDA侧、会释放GIL但调度、切分、汇聚这些操作依然要过主进程。结果就是主卡device 0承担了额外的参数聚合和广播任务显存占用明显高于其他卡同时主卡的通信和计算被串行化成为整个训练流程的瓶颈。我实测过一组数据在一台4卡机器上跑同一个ResNet50分类任务方案单步耗时主卡显存副卡显存GPU平均利用率单卡 baseline0.42s9.8G-92%DP 4卡0.31s12.4G8.6G58%DDP 4卡0.14s9.9G9.9G89%数字很直白DP相比单卡只快了约26%而DDP快了约3倍。主卡显存多吃了2.6G副卡利用率只有58%。这不是个例是DP架构决定的。2.2 DP的核心代码和参数陷阱DP的代码长得像这样import torch import torch.nn as nn model MyNet().cuda() model nn.DataParallel(model, device_ids[0, 1, 2, 3]) for images, labels in dataloader: images images.cuda() labels labels.cuda() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step()看着没有任何问题但有几个隐性地雷BatchNorm的统计问题DP是把每个卡的子batch各自过一遍BN然后PyTorch的DataParallel会把各卡的BN统计通过一个隐式同步机制合并。听起来合理但同步本身是额外通信且如果某张卡的子batch里样本数太少比如batch size8、4张卡每卡只有2个样本BN统计会严重失真训练不稳定。device_ids与可见卡不一致如果你设置了CUDA_VISIBLE_DEVICES2,3但device_ids[0,1]PyTorch会在逻辑卡和物理卡之间产生混乱报错还算好的有时候会静默跑到错误的卡上。稳妥写法是device_idslist(range(torch.cuda.device_count()))。随机数一致性DP下各卡的随机种子由主进程统一控制做数据增强时可能出现不同卡取到相同随机数的情况。DDP下每个rank独立设种子反而更可控。注意DP在较新的PyTorch版本中已经进入了低维护状态官方文档明确建议单机多卡优先使用DDP。如果你的项目还在用DP短期内能跑就先跑但新项目不要再选它。2.3 DP的典型症状和自查方法用DP跑一段时间后如果出现这些症状基本可以确诊nvidia-smi里device 0的显存占用比其他卡高出一截训练速度随卡数增加收益急剧衰减甚至4卡不如2卡某个副卡的utilization长期低于20%反向传播后有一段时间所有卡都在等主卡自查最直接的方式是在训练循环里插一段打印import time start time.time() outputs model(images) forward_time time.time() - start # 对比 forward 和 backward 的耗时占比如果backward耗时占比明显偏高比如超过forward的2倍且卡数越多越明显那就是梯度聚合阶段的主卡瓶颈。这种情况下代码层面的调优空间已经很小换DDP是正解。3. DDP的原理拆解多进程、多通信、无主卡瓶颈DDP全称torch.nn.parallel.DistributedDataParallel它的设计哲学和DP完全相反每张卡一个独立进程各自维护完整的训练状态通过集合通信collective communication在进程间同步梯度。没有主卡概念没有GIL限制没有单点瓶颈。3.1 进程模型和初始化spawn与init_process_groupDDP的第一步是初始化进程组这是新手最容易翻车的地方。核心代码是import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): os.environ[MASTER_ADDR] localhost os.environ[MASTER_PORT] 12355 dist.init_process_group(backendnccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) def cleanup(): dist.destroy_process_group()这里几个参数的意义值得掰开讲backendncclGPU训练必须用NCCL。NCCL是NVIDIA针对多GPU通信优化的库支持NVLink和PCIe拓扑感知的ring all-reduce。CPU训练可以用gloo但GPU上用gloo会慢很多。MASTER_ADDR和MASTER_PORT进程组需要一个集合点来协调各进程。单机多卡的时候填localhost和一个空闲端口即可多机时填rank 0所在机器的IP。rank和world_sizerank是当前进程的全局编号world_size是总进程数。注意rank是全局的不是每台机器内部从0开始。两台机器各4卡world_size8机器A的rank是0~3机器B的rank是4~7。torch.cuda.set_device(rank)让当前进程绑定到对应的卡。这一步非常关键不做的话所有进程默认用device 0直接OOM。然后用torch.multiprocessing.spawn启动多个进程def main(): world_size torch.cuda.device_count() torch.multiprocessing.spawn( train, args(world_size,), nprocsworld_size, joinTrue ) def train(rank, world_size): setup(rank, world_size) model MyNet().to(rank) model DDP(model, device_ids[rank]) # ... 训练循环3.2 梯度同步all-reduce到底在算什么DDP的核心是每一步反向传播后的all-reduce操作。这里要讲清楚梯度平均和梯度求和的区别很多人会混淆。标准的数据并行等价于单卡大batch训练数学上要求对多卡梯度取平均。但DDP默认做的是求和all-reduce在NCCL里默认是sum然后在优化器更新时通过一个系数来等效。更准确地说DDP在初始化时会为每个参数打上bucket在反向传播过程中梯度按bucket触发all-reduce反向后自动完成同步你不需要手动调用任何同步代码。这就带来一个重要的细节DDP的all-reduce是梯度就绪即触发的。因为反向传播是从最后一层往前逐层计算的DDP会把梯度按bucket打包某个bucket的梯度全部就绪就立刻发起all-reduce与后续层的反向计算重叠overlap。这就是DDP比DP快的深层原因之一——通信和计算并行进行而不是等所有梯度算完再统一同步。要理解通信量可以做个估算。假设模型参数量为P以float32存储每个参数4字节一次all-reduce对应的通信数据量约等于2Pring all-reduce的通信量为2(N-1)/N × PN为卡数近似2P。以1.2亿参数的模型为例每次all-reduce要传输约960MB数据。如果PCIe带宽是20GB/s理论耗时约48ms。如果单步计算耗时100ms那通信开销占比接近50%这就是为什么大模型训练非常依赖NVLink和高速网络。提示NCCL的ring all-reduce在卡数较大时通信量与卡数基本无关趋于2P这是它比朴素broadcast-gather方案高效的关键。选型时如果看到gloo更省事别信GPU场景下gloo通常是NCCL的数倍慢。3.3 DistributedSampler数据切分不能想当然DDP下每张卡是一个独立进程如果直接用普通DataLoader每个进程都会独立shuffle同一个数据集导致各卡看到的样本有大量重复等价于缩小了有效batch并且每个epoch里总样本数被重复计算了world_size倍。正确的做法是用DistributedSamplerfrom torch.utils.data.distributed import DistributedSampler sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue, drop_lastTrue ) dataloader DataLoader( dataset, batch_sizeper_gpu_batch, samplersampler, num_workers8, pin_memoryTrue )这里有几个实操要点每个epoch开始前必须调用sampler.set_epoch(epoch)否则shuffle的随机种子每轮都一样数据顺序固定会显著影响收敛。这个坑我第一次用DDP时踩过loss下降曲线正常但最终精度比单卡低1个多点排查了很久才发现是忘了set_epoch。drop_lastTrue强烈建议开启。如果最后一批样本数不能整除world_size各卡的batch大小会不一致all-reduce时出现形状不匹配的报错或隐式pad得不偿失。num_workers要合理设置。每个进程的DataLoader都是独立的4卡×8 workers就是32个进程在读取数据IO压力会叠加。建议num_workers 单卡时的一半到全量观察CPU负载调整。pin_memoryTrue配合.cuda(non_blockingTrue)能减少H2D拷贝的等待时间小改动收益明显。4. 从DP迁移到DDP一个可复现的完整示例纸上谈兵不够我给一个从零可跑的完整流程用的是最简单的MNIST或CIFAR思路方便你在自己的环境里验证。核心结构换成任何模型都能套用。4.1 环境准备和依赖确认先确认环境干净python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.device_count())输出应该能看到版本号、True和你的卡数。如果cuda.is_available()是False别急着写DDP先解决CUDA驱动和PyTorch版本匹配问题——并行训练的错一半以上源于环境没对齐。常见的环境对齐清单组件检查方式版本匹配原则显卡驱动nvidia-smi驱动版本需支持所装CUDA的版本CUDA Runtimenvcc --version或 torch的torch.version.cuda需与驱动兼容PyTorch自带runtimePyTorchtorch.__version__与CUDA runtime对应如cu121/cu118NCCLtorch.cuda.nccl.version()随PyTorch分发一般无需单独装4.2 完整训练脚本模板下面这个脚本在单机多卡上可以直接跑四卡就是nprocs4import os import torch import torch.nn as nn import torch.optim as optim import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from torchvision import datasets, transforms def setup(rank, world_size): os.environ[MASTER_ADDR] localhost os.environ[MASTER_PORT] 12355 dist.init_process_group(nccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) def cleanup(): dist.destroy_process_group() def build_model(): return nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64 * 8 * 8, 10) ) def train(rank, world_size, epochs5): setup(rank, world_size) torch.manual_seed(42 rank) transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) dataset datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) sampler DistributedSampler(dataset, num_replicasworld_size, rankrank, shuffleTrue, drop_lastTrue) loader DataLoader(dataset, batch_size64, samplersampler, num_workers4, pin_memoryTrue) model build_model().to(rank) model DDP(model, device_ids[rank]) criterion nn.CrossEntropyLoss().to(rank) optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9) for epoch in range(epochs): sampler.set_epoch(epoch) # 关键每轮重置shuffle model.train() total_loss 0.0 for images, labels in loader: images images.to(rank, non_blockingTrue) labels labels.to(rank, non_blockingTrue) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() if rank 0: print(fepoch {epoch} loss {total_loss / len(loader):.4f}) cleanup() def main(): world_size torch.cuda.device_count() torch.multiprocessing.spawn( train, args(world_size,), nprocsworld_size, joinTrue ) if __name__ __main__: main()4.3 学习率、batch size与warmup的联动调整DDP下最容易忽略的一环是超参要跟着有效batch size走。假设单卡batch size是324卡就是128的有效batch。线性缩放规则linear scaling rule建议学习率也相应放大但直接乘4往往会震荡实操中更稳的做法是lr_multi_gpu lr_single * sqrt(world_size) # 保守缩放 或 lr_multi_gpu lr_single * world_size # 激进缩放需配warmup warmup_steps 5 * steps_per_epoch # 约5个epoch的warmup我在实际项目里更偏向保守缩放加warmup的组合。直接按卡数线性放大、又没有warmup前几百步loss经常炸。另外梯度累积gradient accumulation和DDP一起用时要注意累加的等效性DDP是每个micro-batch后自动all-reduce如果用梯度累积想模拟更大batch需要通过model.no_sync()包住中间几步只在最后一次同步否则等于白做。这点后面问题排查里会展开。5. 常见问题与排查技巧实录这部分是我认为整篇里最值钱的内容都是真金白银换来的。DDP的报错信息往往指向模糊得靠经验缩小范围。5.1 十大高频问题速查表现象可能原因排查动作解决方向所有进程都占了device 0未调用torch.cuda.set_device(rank)打印torch.cuda.current_device()在setup里显式set_deviceall-reduce挂起不返回rank间通信不匹配某个rank提前退出检查各rank日志是否同时到达同步点统一各分支的集合通信调用次数报错Expected to have finished reduction in the prior iteration某些参数没用上未被反向检查是否有未参与loss计算的模块设置find_unused_parametersTrue训练loss正常但精度比单卡低忘调sampler.set_epoch或用普通DataLoader对比数据生成顺序换DistributedSampler并set_epochBN层统计异常、loss波动大每卡batch太小看每卡batch大小每卡batch≥16或换SyncBN显存比单卡明显增加NCCL缓冲区占用查nvidia-smi里NCCL占用设NCCL_BUFFSIZE或减少bucket端口被占用启动失败MASTER_PORT冲突netstat查端口换空闲端口多机训练卡住防火墙/网络不通从机ping主机MASTER_PORT开放端口、检查网卡速度提升不达预期通信瓶颈或数据加载瓶颈profile单步耗时构成加NVLink、调num_workers进程崩溃后残留进程占卡无join或异常未cleanupps -ef查残留清理进程、加try-finally5.2 三个真实排查案例案例一all-reduce永久挂起。有次训练跑到第2个epoch就卡住不动nvidia-smi显示所有卡都在但利用率100%不动。查了半天发现代码里有个条件分支——某张卡上有个样本触发了某种异常处理走了一条不执行loss.backward()的路径导致那张卡没进入all-reduce其他卡在等它全局死锁。教训所有rank进入的集合通信次数必须严格一致任何rank独有的分支包括异常处理、early stopping都要保证通信对齐。案例二find_unused_parameters的代价。有个多任务模型部分head在某些batch上没有监督信号导致反向时对应参数没有梯度DDP报expected to finish reduction。有人建议直接开find_unused_parametersTrue。开了确实不报错了但速度掉了约20%。原因是这个选项会让DDP每步遍历计算图找未使用的参数额外开销不小。我的处理方式是在模型里把不参与当次loss的模块显式requires_gradFalse或者用多个优化器和多组DDP包装按需激活比全局开这个开关更划算。案例三梯度累积与DDP的冲突。想用大batch但显存不够于是做梯度累积累积4个micro-batch再optimizer.step()。代码写完后发现loss曲线和预期不符。原因就是每个micro-batch的backward都触发了DDP的all-reduce等于在累积阶段同步了4次且每次只同步了1/4的梯度大batch的效果没模拟出来反而多付了通信。正确写法是用上下文管理器for i, (images, labels) in enumerate(loader): with model.no_sync() if i % accum_steps ! 0 else nullcontext(): outputs model(images) loss criterion(outputs, labels) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()model.no_sync()的有效范围是跳过本轮all-reduce只在最后一次同步配合损失除以累积步数才能等效大batch。5.3 性能调优的额外技巧除了上面这些还有几个常被忽视的点model DDP(model, device_ids[rank], output_devicerank)在单机多卡时指定output_device能避免输出被默认放到device 0。梯度bucket调优DDP有个bucket_cap_mb参数默认25MB。通信小块太多会增大开销大块太少又不利于overlap。通信密集的场景可以试着调到50或100观察单步耗时变化。BN替换为SyncBN如果每卡batch实在开不大比如检测任务把nn.BatchNorm2d换为nn.SyncBatchNorm它会把跨卡的BN统计同步效果接近大batch。代价是额外的all-reduce通信。用torchrun启动更省心torchrun --nproc_per_node4 train.py会自动设置rank和world_size等环境变量代码里改成从环境变量读取比手写spawn更适合多机和生产环境。对应的读取方式是local_rank int(os.environ[LOCAL_RANK])。日志只在rank 0打多进程下所有rank都print会刷屏用if rank 0包住输出或者改用torch.distributed的日志工具。我个人在项目里的体会是DP到DDP的切换代码量并不大真正的成本在于理解通信模型和养成每个rank行为要对齐的思维习惯。第一次改成DDP时光排查find_unused_parameters的问题就花了大半天但一旦这套心智建立起来后面再上多机、再上更复杂的并行策略都是顺水推舟的事。小技巧调试DDP时把world_size先设成1跑通再开到2、4问题定位会清晰很多。
