PyTorch单GPU到多GPU训练迁移：DataParallel与DDP的踩坑对比

做深度学习训练最经典的一条路就是从单卡开始跑通再往多卡上搬。PyTorch单GPU和多GPU训练表面上看只差一个“并行”的壳但真正动手改代码时会发现一堆细节需要重新理解。我自己在单卡上跑得好好的脚本第一次用DataParallel包一下之后出现过loss曲线异常、主卡显存爆掉、多进程下随机数全部重复的问题。后来换成DistributedDataParallel又把BatchNorm、学习率、DataLoader这些环节逐个踩了一遍坑。这篇内容不写长篇大论而是把单GPU与多GPU训练之间那些容易忽略的细节做一个对比帮你在迁移之前心里有数。无论你是刚跑通第一个模型还是已经在单卡上训练了很久打算提速这篇文章都能帮你少踩几个真正的坑。1. 单卡训练所有对比的基础1.1 流程和容易出错的位置单GPU训练看起来太基础但它是后续一切对比的基准。标准流程一般分为数据准备、模型构建、损失与优化器、训练循环。我在刚入门时一个最常见的错误是忘了把一批张量统一放到GPU上。模型的参数在cuda上输入数据还在CPU上或者某个中间变量在CPU上代码不会报错因为PyTorch会在某些操作里自动尝试隐式转换但一旦碰到多个输入混合就容易直接报Tensor for argument mat2 is on CPU, but expected ... CUDA。在单卡训练时需要关注的细节其实不多。模型调用.to(device)输入和标签通过列表推导放到同样设备损失函数里如果有额外weight参数也要记得放device。输出logits和label不要一个在cuda一个在cpu交叉熵会给你明确报错。反向传播之前一定要optimizer.zero_grad()否则梯度累积会让人摸不着头脑。虽然这些基础操作看起来简单但多卡场景下很多问题就是从单卡中被忽略的设备归属和生命周期里蔓延出来的。此外我习惯在训练前写一个快速sanity check取一个小batch把模型跑一遍看看前向和反向能不能完整执行。这个习惯从单卡延续到多卡能帮你把“模型代码问题”和“并行框架问题”分开。因为在多卡环境里如果模型本身有bug报错信息会被包装得特别复杂很难一眼锁定根因。先在单卡确认无误再迁移多卡能省下大量时间。1.2 设备放置和DataLoader的细节单卡训练的设备放置一般有两种写法。一种是直接写过死的cuda:0另一种是写device torch.device(cuda if torch.cuda.is_available() else cpu)。我推荐第二种因为迁移到多卡时每张GPU的编号不是固定好的尤其在使用多进程时必须让每个进程绑定不同的GPU。如果你在代码里写死cuda:0多卡DDP环境下八个进程都会去抢0号卡程序能跑但速度奇慢甚至会OOM。DataLoader在单卡上的参数往往被低估。num_workers0是默认很多人图省事不改结果CPU利用率低数据加载跟不上GPU。我一般用num_workers4或8pin_memoryTrue。pin_memory会让DataLoader使用的Tensor锁页之后从CPU到GPU的传输更快。但我不想一开始就建议盲目调大因为num_workers是进程数每个进程都会复制一份数据内存占用会随之涨。单卡时4-8个worker足够多卡时需要考虑总进程数这个我放到后面详细展开。单卡训练还有个重要细节当使用 DataLoader 时如果num_workers 0主进程和数据加载进程之间通过IPC传递数据如果你在collate_fn里写了一些不可pickle的逻辑会报cant pickle。多卡环境下这种报错更容易出现因为DDP多进程本身就是fork出来的建议在collate_fn以及Dataset里尽量只保留简单的可序列化逻辑。2. 多GPU训练的两条主流路线DataParallel vs DistributedDataParallel2.1 DataParallel的便捷与隐患最省事的多卡方案是 DataParallelDP。只需要一行model torch.nn.DataParallel(model)原来的单卡代码几乎不用动。DP在机制上是把输入batch在第0张卡拆成多个sub-batch分发给各个GPU计算然后由0号卡汇总梯度、更新参数。听起来很方便但实际使用中有几个很明显的坑。坑之一主卡显存负担重。因为输入数据、梯度合并、参数更新都在0号卡上做所以0号卡的显存占用会明显高于其他卡。我原来用2张卡训练每张卡12G显存batch size稍微大一点0号卡就先OOM另外一张卡还只用了不到一半负载极不均衡。这个不是数据排布问题而是DP架构上的天然缺陷。坑之二进程数并没有变。DP只是把多GPU放在一个进程内靠多线程方式并行。我记得PyTorch在GPU计算时通过CUDA stream可以实现并行但python的反向传播、数据收集等逻辑仍然受GIL影响实际加速比往往达不到卡数倍甚至有些模型只提升了30%-50%。多机场景更没法用DP只适用单机多卡。坑之三BatchNorm在每个GPU上独立统计。DP把batch分到多卡每张卡上的子batch可能很小。假如全局batch size是322卡每卡只有16本来BatchNorm按32算均值方差现在变成了按16算甚至2卡变8卡后每卡只有4统计值就不稳定训练收敛会变得更差。所以用DP时要小心BN层较多的网络。2.2 DistributedDataParallel的正确起飞姿势相比DP更推荐的是 DistributedDataParallelDDP。DDP的关键思路是每个进程拥有一份完整模型副本进程数等于GPU数量每个进程只读取自己那一份数据前向反向各自独立反向传播时通过通信操作将梯度在进程间进行allreduce同步然后每个进程各自更新参数。我第一次见识DDP的速度优势时才知道之前用DP纯粹是在跟自己过不去。DDP的同步方式对通信整体做了梯度分桶bucket不是每个参数单独通信这样减少了大量小消息的延迟。初始化步骤一般像这样import torch.distributed as dist dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(local_rank) model model.to(local_rank) model nn.parallel.DistributedDataParallel(model, device_ids[local_rank])然后要用 torchrun 进行多进程启动而不是直接python train.py。torchrun 会自动设置 RANK、LOCAL_RANK、WORLD_SIZE 这组环境变量。数据加载部分要配合 DistributedSampler让每个进程拿到不同的数据切片否则多个进程会读重复数据。这里有几个容易犯的错误只初始化了进程组但忘了 set_device只给模型套DDP但忘了用 DistributedSampler或者一个进程里同时用了DP和DDP造成状态混乱。一个常被忽略的细节是DDP默认会通过rank0的模型参数和buffer广播让所有进程初始状态一致。整个过程发生在构造DDP模型时不需要手动写广播代码。所以如果你在构造模型之后修改了模型内容记得重新包装或者手动同步否则各个rank的参数会不一致。DDP的前向反向逻辑和单卡几乎一致但多了一个“梯度同步”环节。默认情况下DDP对梯度做的是平均而不是求和具体取决于PyTorch版本和设定但常见理解是allreduce后除以world_size得到平均梯度。有一个细节每个进程的损失是各自的batch loss如果是平均loss全局的loss含义和单卡不太一样。我们在记录训练日志时最好统计所有进程loss的平均而不是只看rank0的loss否则数值可能有比较大的偏差。为了快速对比下面用一个表总结DP和DDP的主要区别对比维度DataParallelDistributedDataParallel进程模型单进程多线程多进程每个GPU一个进程适用场景单机多卡单机多卡、多机多卡同步方式主卡汇总梯度allreduce跨进程梯度同步显存负载主卡负载高相对均衡数据加载所有数据经主卡分发各进程独立加载子集推荐程度适合临时验证正式训练推荐2.3 为什么我建议从单卡直接迁移到DDP很多人觉得DP代码改动小先用DP做基线再换DDP。但我实际经验是DP几乎只能在命令快速验证时用一旦涉及正式训练各种问题会接连出现。不如从一开始就按DDP的规范来写这样代码的扩展性更好。尤其是当项目后续要跑多机或者需要接入SyncBatchNormDP反而是技术债。DDP的启动方式也不复杂。用torchrun跑一遍把训练脚本的参数带上即可。虽然代码改动比DP多但模式很固定学一次就够。而且DDP在很多官方模型库中已经成为默认方案你搜到的开源项目多数也是DDP。跟随社区的主流做法后续排查问题也更容易找到同类案例。3. 训练策略细节对照Batch Size、学习率、BatchNorm和随机种子3.1 Batch Size改变之后学习率要不要动单卡训练的batch size通常受显存限制。多卡后如果每张卡上仍然保持原来的batch size那全局batch size就是原来的N倍。global_batch_size per_gpu_batch_size * world_size。如果只是把模型从单卡换成多卡但不改变总样本数那么一个epoch里的优化器更新步数会减少到原来的1/N。这会对训练产生微妙影响。学习率需要跟着调整吗常见做法之一是线性缩放规则当batch size乘以N时学习率也乘以N。因为更大的batch size对应更平滑的梯度估计可以采用更大的学习率但也不是无脑乘通常还需要配合warmup。我在实际项目中从batch 32升到batch 64时学习率从0.001调到0.002训练曲线的收敛速度能匹配上原来的节奏。有些任务的batch size本来就够大或者模型结构偏好小学习率那就得多做几次实验。反过来也有人会让每张卡的batch size变小以保持全局batch size不变。例如单卡batch644卡时每卡16全局还是64。这样参数更新频率和单卡一样但每卡上的BatchNorm统计量变成了16模型中的BN会有较大波动。所以通常不建议简单地把全局batch size保持不变而是尽量保持per-GPU batch size同时对学习率和训练步数做适配。还有一个细节是梯度累积。有时我们想让batch size更大但显存撑不住就用梯度累积模拟。此时global batch size per_gpu_batch_size * world_size * accumulation_steps。学习率该如何设置就没有一个简单公式要看累计等效batch和真实batch的关系。我建议在设定累积步数时先固定per-GPU batch不要太小再用累积把全局batch补到目标值学习率按全局batch来线性调整但每积累一次loss后实际上梯度是多个batch的平均有时候步数多了要按sqrt scaling或者保持常数。这些都要通过小实验验证。3.2 BatchNorm在多卡环境下的身份微妙BatchNorm的原理是在一个mini-batch内计算均值和方差。单GPU时这个“batch”就是当前这一个设备的batch。多卡DP/DDP默认情况下每张卡各自计算各自的均值和方差并不和别的卡通信。如果你的全局batch size很大但每张卡的batch size很小BN的统计量会严重偏离全局影响收敛。这就是为什么有同步BatchNormSyncBatchNorm这个东西。它让所有GPU协作先通过通信把各设备上的统计量汇总再计算出全局的均值和方差等效于在更大batch上做BN。在PyTorch里对DDP包装后的模型可以调用torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)把已有的BatchNorm层替换成SyncBatchNorm实现。但SyncBN的劣势也很明显每次前向BN统计都需要跨进程通信训练速度会下降如果每个GPU的batch已经足够大比如32以上通常不需要SyncBN。实践中我用过一张卡batch4的情况模型带BN单卡训练时效果还正常但改成4卡DDP后每张卡batch4发现验证集的表现明显下降。后来检查发现是因为每卡BN只统计4个样本噪声太大。改成SyncBN后等效batch16效果才回升。但要注意如果模型同时使用了GroupNorm或者LayerNorm它们不受batch影响不需要同步。还有一个相关坑有些模型的BatchNorm被设为track_running_statsFalse也就是用batch统计而不用running stats。在多卡下这种BN的统计量天然每卡不同迁移到测试时可能需要特殊处理。如果你的模型里包含这种结构请留意自己在训练和验证时使用的模式尽量保持一致。3.3 随机种子在多进程里容易让人迷路很多同学会在代码里设置一个随机种子比如random.seed(0); np.random.seed(0); torch.manual_seed(0)但搬到DDP多进程后如果不加处理每个进程都会使用相同的随机种子。这会导致什么数据加载的shuffle、dropout等随机行为在多个进程里完全一致最终所有GPU都在看同一份数据的相同顺序等于一个模型重复训练了N份。这样加速完全没有意义因为全局的有效batch并不是N倍而是只有原来一份。正确做法是让每个进程的随机序列不一样。我自己的习惯是base_seed 42然后每个进程设置torch.manual_seed(base_seed rank)。同时针对DataLoader的DistributedSampler它内部有一个epoch计数器在不同epoch里会进行不同的shuffle。但如果你不记得在每轮迭代开始时调用sampler.set_epoch(epoch)那么每个epoch的shuffle顺序会保持不变尽管随机种子不同但sampler内部用同样的顺序仍然会让数据分配不够多样化。所以DDP的标准代码是sampler DistributedSampler(dataset) train_loader DataLoader(dataset, batch_sizebatch_size, samplersampler) for epoch in range(epochs): sampler.set_epoch(epoch) # 正常训练循环另外在模型构造阶段如果有batch_norm的affine参数、dropout等随机性也需要在构造前把基础seed设置为seed rank。我记得在一些老版本PyTorch上DDP的rank 0会广播初始模型参数就算各个进程seed不同模型初始参数也一样因为最后是rank0的参数同步过去。但这不影响训练中的数据shuffle和dropout。4. 资源与性能的细节调优4.1 DataLoader的worker数量和pin_memory千万别翻倍后就忘了单卡DataLoader设置num_workers8可能没问题。多卡DDP时如果世界大小是8那么每个进程都会创建一个拥有8个worker的DataLoader总数就是64个worker系统会瞬间拉起大量进程内存占用飙升CPU频繁切线程数据加载反而变慢还会增加OOM概率。我的经验是每张GPU对应的进程DataLoader的worker数不需要那么大。在本地磁盘读取图像数据时单卡num_workers8可能比较合适但8卡DDP时每个进程设置num_workers2或3总worker约等于16-24已经能很好地填满IO带宽。多卡加速的核心是并行计算数据加载如果变成瓶颈先看总worker数不要只看单进程配置。也可以在DataLoader里加上persistent_workersTrue避免每个epoch结束后杀掉worker再重新创建的开销。但persistent_workers要求迭代完一定不能忘干净配合合适的shuffle策略还好。pin_memory在多卡时要特别注意。pin_memoryTrue会在主机内存中分配锁页缓冲区将数据从pageable memory转到pinned memory再拷贝进GPU会更快。多进程时每个进程都有一份锁页内存显存和内存的压力都会增加。如果内存小建议只给部分DataLoader开或者调低worker。另一个小技巧是如果你用DDP每个进程只负责读自己的子集而不是像DP那样所有数据都先load到主进程所以每个DataLoader的总数据量是原来的1/N实际上IO压力小了很多worker数就不需要那么大。还有一个细节是collate_fn。多卡场景下不要在每个进程里把同一份大文件重复读到内存。比如有些人不小心在Dataset__init__里把所有数据全部load成一个list多进程会各自复制一份内存直接被爆。数据读取尽量采用惰性在__getitem__里做并且通过文件路径读取避免重复embedding整个数组。4.2 混合精度与梯度累积在多卡下的配合细节混合精度训练已经成了标配尤其在大模型和单卡大batch时能显著提速。PyTorch提供了torch.cuda.amp.autocast和GradScaler单卡的时候用法简单scaler torch.cuda.amp.GradScaler() for batch in loader: optimizer.zero_grad() with torch.cuda.amp.autocast(): loss model(inputs) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()多卡DDP下建议每个进程各维护一个scaler实例因为loss是各进程自己算的梯度同步发生在allreduce之前。理论上同一模型中每张卡的参数完全一样每张卡的连续两次step之间的scaler状态可能略有差异但差异不大一般不用手动同步。我在AMPDDP里遇到过一个坑如果过度依赖scaler内部判断如果某一卡出现inf梯度其他卡可能受影响这时需要检查输入数据和loss是否合理而不是简单关闭AMP。梯度累积在多卡场景下特别常见。DDP会同步每个mini-step的梯度如果你通过累积梯度来等效增大batch实际上每次backward后都会触发一次梯度同步。虽然这一步看起来是白跑的通信但为了保持模型状态一致必须这么做。要注意的是在累积模式下optimizer.step()只在满累积步数时执行但每次backward后仍然会进行allreduce因此通信开销不变。有人会用no_sync()上下文管理器跳过部分累积过程中的梯度同步只在最后一个mini-step执行同步。PyTorch提供了model.no_sync()上下文在这个上下文里不会同步梯度。但要注意梯度仍然是累加的最后一步退出上下文时完成allreduce。这个技巧能显著降低通信量但是要求数据拆分策略和loss计算方式能够保证数学等价否则会导致梯度平均的语义混乱。另外BN和AMP也有一个配合注意点如果把模型转成SyncBN后再用autocastSyncBN内部的通信在fp16下可能会产生精度问题PyTorch里的SyncBN转换时一般会保留某种精度但最好在把模型放进DDP之前完成转换并且不要在autocast里包模型的前向之前使用高精度BN。我建议不搞太花哨先按官方推荐顺序model torch.nn.SyncBatchNorm.convert_sync_batchnorm(model); ddp_model DDP(model);然后在forward外面用autocast。实测这套顺序比较稳。5. 常见问题排查实录5.1 module前缀和加载权重时的坑将单卡模型放到DP或DDP之后state_dict的key会发生改变最典型的是原始键比如backbone.conv1.weight变成了module.backbone.conv1.weight。保存模型时如果你直接保存ddp_model.state_dict()加载到普通单卡模型时会报unexpected key module.backbone...然后模型参数还是随机初始化很多人在这里栽跟头。我常用的保存姿势是在模型保存之前先取出原始模块。如果是DP用model.module.state_dict()如果是DDP也是如此。但有时为了稳妥不管里面有没有module前缀加载时都可以做一个兼容处理state_dict torch.load(path, map_locationcpu) if any(k.startswith(module.) for k in state_dict): state_dict {k[len(module.):]: v for k, v in state_dict.items()}这个技巧对单卡加载多卡权重很实用。另一个细节是DDP模型里除了普通参数还会有一些buffer比如BatchNorm的running_mean、running_var。多卡保存时每个进程的buffer理论上一致但保险起见保存参数会使用rank0的模型。加载时map_locationcpu可以避免当前GPU编号不一致的问题。如果用了SyncBatchNorm权重key同样会多出module前缀但内部实现可能有auxiliary状态比如num_batches_tracked之类的兼容处理时不要漏掉。最简单的方案是保存时只存model.module加载时也只用原始模型结构加载再在需要时重新包装。5.2 DDP卡死和超时问题DDP训练卡死是我最头痛的问题之一。典型表现命令行卡住没有任何输出或者只有初始化信息然后进程不退出。原因通常有几种某个进程因为数据不足先结束而其他进程还在等它同步DataLoader的worker数或者数据拆分不均匀导致某个rank没有batch还有可能是初始化端口冲突或者网络通信问题。最常见的一个是数据量不能被world_size整除。DistributedSampler默认是drop_lastFalse它会填充一些样本使得每个rank样本数一致。但如果你用了别的自定义sampler没考虑整除就可能出现rank0跑了10个batchrank1只跑了9个batch等第10次allreduce时rank1迟迟不调用于是rank0卡死等待。解决方式要么确保数据集长度是world_size × per_rank_batch的整数倍要么在sampler上设置drop_lastTrue要么用DistributedSampler自带的padding。另一个是init_method。如果使用env://确保每个进程的环境变量都有RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT。用torchrun时这些变量会自动注入但如果你用终端手动启动多进程很容易漏掉。MASTER_PORT如果和本机已有端口冲突也会初始化失败。我遇到过用5000端口被其他服务占用进程组一直起不来。排查时用python -m torch.distributed.run --nproc_per_node... train.py会比手动spawn更省心。timeout参数也很关键。dist.init_process_group里有timeoutdatetime.timedelta(seconds1800)的配置默认是1800秒不代表1800秒后会自动卡死。但在分布式训练时如果通信确实由于网络问题一直等不到超时后会有报错而不是死等。我之前在一个共享集群上因为节点间网卡速率低一次allreduce花了超过默认timeout训练直接崩了。把timeout设大一点比如3600能避免虚惊。5.3 显存不均衡和负载不均用DDP时显存一般比较均衡但偶尔会因为某一块GPU上额外执行了验证或日志操作导致其显存占用高一些。排查方法很朴素在训练中途用GPU监控命令查看每张卡的使用率。我遇到过一个案例两张卡的使用率始终是70%和30%后来发现代码里把验证集也放到了rank0的GPU上而且日志和评估每轮都在rank0跑导致那张卡严重拥塞。解决办法是让所有评估和日志只在rank0做但把临时tensor及时del并清一下缓存其实验证也可以放在CPU或者专门留一张卡但最省事的办法是评估时使用torch.no_grad()并将验证模型参数拷贝到CPU上做配合torch.cuda.empty_cache()虽然不能根治瓶颈但至少不会因为显存导致OOM。还有一个隐蔽问题当模型输出非常大比如每张卡输出一个巨型logits然后你去gather所有rank的输出很容易导致rank0显存暴涨。分布式训练中我们有时候需要保存所有样本的预测结果用torch.distributed.all_gather会很方便但all_gather会在每个rank上都持有完整数据等于显存乘以world_size。在大输出量场景下建议改用分批写入文件或者使用all_gather_object存文件路径别把所有大张量集中在一个GPU上。我踩过一次rank0爆显存其他卡很空就是因为all_gather把所有层的特征图都收回来了。负载均衡还有另一个层面模型的某些层计算量不均匀比如第一张卡包含embedding层参数量巨大但计算量占比不高可能导致不同卡的耗时并不完全一致。DDP的同步步调取决于最慢的进程所以即便是DDP提速也可能达不到线性。非线性的提升有时候是正常的不要以为是代码错了。可以通过torch.profiler查看各rank的耗时分布帮助定位。不过对于一般业务重点是别让某一卡塞了额外任务影响整体效率。5.4 一个关于梯度同步方向的细节提醒有一个细节在排查bug时很容易忽略DDP的梯度同步是基于bucket的。默认bucket_size_mb25如果模型参数不多一个bucket就能涵盖所有参数同步效率高。但如果模型非常大可能会分成多个bucket每个bucket的同步时机不同。代码中如果尝试在backward之后、optimizer.step之前读取某个参数的梯度可能在部分rank上已经更新部分rank上还没同步。这时你会看到梯度值不一致但其实DDP设计如此。标准做法是在optimizer.step之后或通过DDP提供的hooks来处理梯度不要手动在backward后直接修改梯度除非你非常清楚时机。从我个人经验来看统一使用DDP是省心又省力的做法。如果你的代码在单卡上已能稳定运行迁移多卡时尽量保持per-GPU batch size不变配合DistributedSampler和torchrun学一下线性缩放规则再花点时间把DataLoader worker和混合精度的顺序理顺就会发现多卡训练并不神秘。后续还可以考虑把模型扩展到更大的数据并行策略但基础细节还是在单GPU与多GPU切换中最值得打磨。希望这些踩坑记录能帮你少走我走过的弯路。