刚把手头一个多卡训练的实验收尾正好有朋友问起PyTorch单GPU和多GPU训练到底差在哪儿、有哪些坑索性把这次踩过的细节一次性梳理出来。这篇文章不会去教你怎么装CUDA、怎么跑通一个分类网络而是聚焦在真正工程化之后绕不开的那些对比和细节数据怎么喂、梯度怎么同步、batch size怎么算、模型怎么存怎么加载以及在单卡上写得好好的代码为什么挪到多卡就翻车。1. 为什么单卡和多卡训练远不止“改几行代码”那么简单先给结论把单卡代码改成多卡表面上看是加几行分布式初始化、包一个DDP真正的问题是整套训练的“行为”变了。你原来熟悉的batch size、学习率、BN统计、随机种子、评估指标在多卡环境下的语义都会发生微妙但致命的变化。单卡训练的逻辑是这样的一个进程一块GPU一份数据一个模型。每次迭代读一个batch前向传播算loss反向传播算梯度优化器更新参数。整个过程是串行且确定性的你设了随机种子基本能复现虽然有些算子仍然不是完全确定性的。多卡训练则引入了“群体协作”的概念。最常用的DataParallelDP和DistributedDataParallelDDP虽然名字差不多但哲学完全不同DP单进程多线程把一块卡当master数据在每个forward之前被scatter到各卡forward各算各的然后gather到master卡算loss和梯度master更新参数后再broadcast给所有卡。DDP多进程每个进程绑定一张卡每张卡都有完整的模型副本数据通过DistributedSampler分片每张卡独立前向、反向然后通过ring-all-reduce算法同步梯度每个进程再各自用同步后的梯度更新自己的模型副本。正因为DDP是“各自为政、梯度共享”它的效率、稳定性、扩展性都远超DP现在已经成了事实上的标准。但问题恰恰出在这里很多从单卡迁过来的人还是用单卡的思维去理解DDP于是踩进一个又一个细节坑。一个非常典型的现象是单卡跑到96%准确率的代码搬到DDP多卡训练后准确率掉了0.5个百分点然后查了半天最后发现是batch size变了、学习率没跟着变或者DistributedSampler没有正确shuffle导致每个epoch的数据分布不随机。这类问题非常隐蔽因为程序不报错loss也在降但就是效果不对。2. 单卡与多卡训练的核心差异拆解2.1 数据加载与batch size语义的变化单卡训练时假设你设置batch_size64那么这个64就是每次迭代喂给模型的实际样本数。多卡DDP训练时如果你仍然设置batch_size64那这里的64是指“每张卡的batch size”全局batch size实际是64 × 卡数。这个差异直接决定了你的梯度估计质量、收敛速度、最终精度。很多人喜欢说“多卡训练效果应该和单卡一样”这句话只在“全局batch size和单卡batch size一致”或者“学习率做了相应缩放”的情况下才成立。具体来说单卡上batch size256跑得稳你搬到4卡DDP想保持全局256就得把每卡batch size设为64。如果你把每卡batch size也设成256那全局就是1024学习率完全不调的话模型大概率从一开始就发散。解决办法是和学习率联动。最常用的经验法则是linear scaling rule当batch size变为原来的k倍时学习率也放大k倍。比如单卡batch256时lr0.14卡全局batch1024时lr应该放大到0.4。但要注意这个规则不是无脑套用的它受batch size上限约束batch太大时梯度噪声已经很小了再线性放大学习率容易震荡所以常用做法是warmup配合缩放比如前几个epoch从很小的lr逐步升到目标lr。除了数值变化还有一个很容易被忽略的点多卡训练里的数据随机性。DDP依靠DistributedSampler来给每个进程分配互不重叠的数据分片。这个sampler会在每个epoch自动进行shuffle吗如果你不手动调用set_epoch()默认情况下每个epoch的分片顺序是完全一样的等于你每个epoch都在用同样的顺序看数据这对训练结果的影响通常不会直接体现在loss上而是体现在最终泛化性能的退化上。2.2 梯度同步机制DDP的all-reduce到底做了什么单卡训练的backward是纯粹的本地计算loss对每层参数求梯度然后优化器拿着梯度更新参数。DDP的backward则多了一步在反向传播的过程中框架自动触发梯度同步。具体机制是DDP在模型注册时会对每个参数注册一个autograd hook。当反向传播进行到某一层时hook会被调用把所有进程上这一层参数的梯度收集起来做一次all-reduce操作。也就是把每个进程的梯度相加或求平均再分发给所有进程。这里有个重要配置torch.distributed.all_reduce的默认op是SUM但DDP实际上做的是“梯度平均”。DDP在构造时会把gradient_as_bucket_view、reduce这些参数处理好默认的bucket是每个25MB一个桶桶内的参数梯度一起reduce减少通信次数。我最早踩过一个坑手动用torch.distributed.all_reduce去同步梯度时忘记除以卡数导致梯度变成原来k倍learning rate不缩放的话loss直接飞了。DDP内部帮你做了平均这种简陋的、自写all-reduce替换梯度同步的方案一定要记得按进程数做平均。从通信开销角度看DDP把梯度同步放在反向传播阶段与计算重叠所以理论上反向传播耗时只比单卡多一点点。真正影响扩展性的是每步迭代必须等待所有进程的梯度都reduce完一旦某张卡因为数据加载慢或者CPU瓶颈拖了后腿全局都会被拖慢。这就是为什么多卡训练一定要用num_workers调好DataLoader的预取并且把pin_memoryTrue打开减少CPU到GPU的拷贝阻塞。2.3 模型保存与加载DDP的state_dict多了一层wrapper这应该是单卡转多卡时最容易破防的报错之一。单卡训练完torch.save(model.state_dict(), model.pth)然后加载时就model.load_state_dict(torch.load(model.pth))一切正常。换成DDP后model被外层包了一层DistributedDataParallel它的state_dict的key全部多了一个前缀module.xxx。所以常见的做法是# 保存时剥离module前缀 torch.save(model.module.state_dict(), model.pth) # 或者加载时剥离key state_dict torch.load(model.pth) new_state_dict {k.replace(module., ): v for k, v in state_dict.items()} model.load_state_dict(new_state_dict)但有更优雅的做法。DDP的state_dict其实提供了两个方法model.state_dict()返回包含module.前缀的完整字典model.module.state_dict()返回原始模型的字典所以保存时用model.module.state_dict()加载时先初始化单卡模型再把权重load进去之后model DDP(model)就完事了。还有一个更隐蔽的坑直接用torch.save(model, model.pth)把整个模型对象保存。在DDP下这个save会把整个分布式训练状态、进程组信息、甚至Rank相关的配置序列化进去换一台机器或者换一个卡数加载时直接报奇怪的pickle错误。我见过不少人在单卡训练成功后就一直用torch.save(model, ...)这个写法搬到DDP后第一次保存也许没报错但二次加载时就各种诡异问题。所以多卡训练环境下务必使用state_dict方式保存。2.4 BN层在多卡环境下的统计量差异默认情况下torch.nn.BatchNorm2d在单卡上直接计算当前batch内样本的均值和方差。DDP下的BN默认行为也是每张卡各自计算当前分片的统计量。当每卡batch size不够大时比如每卡只有8张图BN的统计量会非常不准噪声极大。你可能会想DDP不是有梯度同步吗BN统计量为什么不一起同步因为BN的均值和方差是在forward过程中计算的不是梯度不会走DDP的all-reduce。这就是为什么多卡训练时经常出现单卡batch32时BN呜呜地收敛多卡每卡batch8、全局实际很大但BN表现反而变差。解决方案有三个增大每卡batch size使BN统计量近似单卡效果但这和“全局batch size不变”的要求冲突。开启torch.nn.SyncBatchNorm它会在forward时同步所有进程的均值和方差。用法很粗暴model torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)然后模型内的BN都会自动替换为SyncBatchNorm。注意SyncBatchNorm的同步通信发生在forward阶段会引入额外的等待和通信开销但换来的是每卡小batch下BN统计量的准确。去掉BN用GroupNorm或LayerNorm替代。在小batch场景下GN稳定性和效果往往比带SyncBN的模型还好而且省去了通信开销。2.5 Loss与评估指标在多卡上是否需要特殊处理常见lossCrossEntropyLoss、MSE、L1等默认在batch维度求均值。在多卡DDP下每张卡独立计算loss然后各自反向传播。这个loss数值是“每卡batch的均值”梯度平均后等价于全局batch的均值这个行为是自洽的。但如果你用了一个自定义loss比如per-sample loss求和后除以一个自定义的全局分母情况就麻烦了。举个例子你写了一个loss (pred - target).sum() / global_n在单卡时一切正常到DDP下每张卡都除以同一个global_n梯度平均之后整体梯度就变成了原来的“每卡求和 × 卡数 / global_n”和单卡语义完全不一致。正确的做法是保持loss计算为“每卡batch内求均值”因为DDP的梯度平均天然把所有卡的均值平均等价于全局均值。如果你想用求和类loss就按卡数缩放或者把所有loss单独reduce到master再处理——但这样会打断梯度同步非常不建议。评估指标同理。多卡时如果每张卡各自算accuracy取平均是可以的但如果有的卡分到了9个样本有的卡分到了10个简单平均就会有偏差。更稳妥的做法是分批计算每个进程的指标最后统一用all_gather汇总数据算全局指标。写代码时这点很烦但指标正确性对论文结果来说就是硬指标省不了。2.6 随机种子与可复现性单卡训练时设置torch.manual_seed(0)加上cudnn.benchmarkFalse和deterministicTrue大概率能复现训练过程。DDP下就没这么简单了。首先是数据加载。DistributedSampler带shuffle时内部会根据epoch和一个随机种子生成索引。如果每个进程都用同一个初始种子那么分到的数据分片是互补且确定的但为了进一步保证均衡DDP框架默认还会结合rank信息做偏移。你需要给每个进程设置不同的全局seed否则你只会在一个看似相同的随机状态里打转。我的做法是def set_seed(seed, rank): random.seed(seed rank) np.random.seed(seed rank) torch.manual_seed(seed rank) torch.cuda.manual_seed_all(seed rank)然后是CUDNN的不确定性。多卡下不同进程的卷积计算本身是确定性的只要算法一致但all-reduce的序列和浮点累加顺序可能因为网络拓扑不同导致微小差异。这个问题通常不影响模型最终性能但如果你做消融对比想要每次跑出来的结果几乎一致那我建议多卡训练还是接受那个非常微小的浮动。3. 实操过程一步不漏地搭建单卡/多卡训练3.1 单卡训练的基线代码结构先放一段我常用的单卡训练核心代码后面多卡版本就是从这里改的。import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset import torchvision # 单卡训练基本结构 def train_single_gpu(model, train_loader, optimizer, criterion, epochs, device): model.to(device) model.train() for epoch in range(epochs): running_loss 0.0 for batch_idx, (images, labels) in enumerate(train_loader): images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) epoch_loss running_loss / len(train_loader.dataset) print(fEpoch {epoch1}/{epochs}, Loss: {epoch_loss:.4f})可以看到data_loader可以很随意batch大小就是你设的值没有任何进程概念。这个版本要改成DDP多卡需要动的点还不少下面逐个展开。3.2 从单卡改造为DDP多卡的完整步骤第一步是初始化进程组。这部分放在main函数的最开头。我最常用nccl后端因为它就是GPU集群上的标准方案import torch.distributed as dist import torch.multiprocessing as mp def init_process(rank, world_size, master_port29500): dist.init_process_group( backendnccl, init_methodtcp://127.0.0.1:{}.format(master_port), rankrank, world_sizeworld_size, ) torch.cuda.set_device(rank)如果你用torchrun或者python -m torch.distributed.launch --nproc_per_node4 train.py启动环境变量LOCAL_RANK、RANK、WORLD_SIZE会自动注入代码可以更简洁import os local_rank int(os.environ[LOCAL_RANK]) dist.init_process_group(backendnccl) torch.cuda.set_device(local_rank)然后模型和优化器基本不变只是加一层DDP包装model model.to(local_rank) model nn.parallel.DistributedDataParallel(model, device_ids[local_rank])第三步改DataLoader。这里有两个关键点DistributedSampler必须用shuffle不能在DataLoader里设置from torch.utils.data.distributed import DistributedSampler train_sampler DistributedSampler(train_dataset) train_loader DataLoader( train_dataset, batch_sizebatch_size_per_gpu, samplertrain_sampler, num_workers4, pin_memoryTrue, )注意DistributedSampler内的shuffle参数有默认为True你不需要在DataLoader里再传shuffleTrue。如果你在DataLoader里也传shuffleTrue会直接报错sampler与shuffle不能同时指定。训练循环里每个epoch开始时要额外调用一次for epoch in range(epochs): train_sampler.set_epoch(epoch) # 正常训练这个set_epoch必须做否则如前面所说每个epoch的数据分片顺序一成不变。3.3 单卡/多卡DataLoader配置差异总结这块我把几个容易踩泥坑的配置项单独拎出来做一张对比表保存到备忘录里都不亏。配置项单卡训练DDP多卡训练batch_size实际喂入模型的样本数每张卡的样本数全局batch batch_size × world_sizesampler不设置DataLoader默认shuffle必须用DistributedSampler且每epoch手动set_epochshuffle参数DataLoader(shuffleTrue)DataLoader里不设置由sampler控制num_workers按需设建议每卡至少4最好8否则CPU会成为瓶颈pin_memory建议True强烈建议TrueMPI场景下能显著减少传输阻塞devicecuda:0 或 cuda通过local_rank动态指定很多教程喜欢在DDP的DataLoader里保留shuffleTrue然后问为什么报错其实就是这个原因。3.4 学习率缩放、warmup与梯度累积的取舍多卡训练的batch size变大后学习率必须缩放。但这里面还有两个操作上的取舍。一个常见组合方案是把warmup显式加入训练循环。我习惯这样写def adjust_learning_rate(optimizer, epoch, warmup_epochs, base_lr, total_epochs, scale_factor): if epoch warmup_epochs: lr base_lr * (epoch 1) / warmup_epochs else: # 余弦退火 progress (epoch - warmup_epochs) / (total_epochs - warmup_epochs) lr base_lr * 0.5 * (1 torch.cos(torch.tensor(progress * torch.pi))) for param_group in optimizer.param_groups: param_group[lr] lr.item()这个函数里的scale_factor就是全局batch size除以单卡基准batch size的比值。实际使用中base_lr lr_single * scale_factor有两种选择一种是在构造优化器时直接把lr乘好另一种是在warmup端点乘好。推荐后者因为warmup阶段其实不该被缩放否则刚起步就天然大学习率风险很大。另一个方案是梯度累积。如果你想在有限的显存下模拟更大的batch可以在每卡batch不变的情况下多累积几步梯度再更新一次。但注意梯度累积会放大“稀疏更新”的时间窗口在多卡通信场景下这一步还会导致本来可以和其他卡的梯度同步重叠的计算被迫拆开整体吞吐反而下降。我一般是“能用更大的batch就用更大batch实在显存卡住才用梯度累积”。3.5 常见多卡启动方式和选择逻辑多卡训练一般有三种启动方式我依次说下各自的特点。torch.distributed.launch老牌的启动器通过--nproc_per_node指定卡数。会设置环境变量RANK、LOCAL_RANK、WORLD_SIZE。torchrun官方推荐的替代工具在PyTorch 2.0以后更稳。也支持--nproc_per_node还额外有容错、动态参与等能力。mp.spawn代码内嵌启动的方式适合写小demo不推荐在大规模训练里用。一个关键区别是RANK和LOCAL_RANK的含义。RANK是全局进程编号LOCAL_RANK是当前机器上的进程编号。如果你只有一台8卡机两者相等如果你有4台8卡机每个进程的RANK是0-31但LOCAL_RANK是0-7循环。torch.cuda.set_device必须用LOCAL_RANK因为进程只操作本机上的可见卡。我自己的建议新项目一律用torchrun。torchrun的使用示例torchrun --nproc_per_node4 --master_port29500 train.py在代码里通过local_rank int(os.environ[LOCAL_RANK])获取设备号。这个方式省心省力网上各种教程也都默认支持。4. 多卡训练常见报错与避坑实测4.1 “AssertionError: DistributedDataParallel is not intended to be used with DataLoader shuffleTrue”这类问题这个报错很好理解DataLoader的shuffle和DistributedSampler冲突。DistributedSampler接管了数据的乱序责任DataLoader只需要根据sampler提供索引取数。但我还遇到过升级版的报错使用IterableDataset时DistributedSampler会不起效果。因为IterableDataset没有__len__没法切片框架没法知道该给每个进程分多少数据。这时需要用torch.utils.data.distributed.DistributedSampler的一个变体或者自己在IterableDataset里按worker和rank做分片。更接地气的办法是别用IterableDataset直接用Dataset和DistributedSampler组合。4.2 “Default process group has not been initialized”的报错很多从单卡迁移的人会在main里直接调用DDP包装模型结果报这个错。原因就是忘了在DDP之前调用dist.init_process_group。正确顺序必须是dist.init_process_group(...) model DDP(model, ...)有些教程会把init_process_group放在import之后但放在函数执行前也没问题关键是必须在一个进程组存在后才能做任何分布式操作。4.3 每张卡的显存占用严重不均正常DDP训练中每张卡的显存占用应该几乎一致。如果你发现某张卡往往是rank 0显存远高于其他卡大概率有两个原因模型参数持久化在master卡上但DDP本来会把模型复制到每张卡不太可能出现整体差异你在代码里做了某种不对称操作比如往rank0的device搬了大tensor、保存了中间结果、或者只在rank0上跑了验证集。排查思路在训练循环里临时打印torch.cuda.memory_summary(local_rank)对比每张卡在同样迭代轮次下的cudaMalloc和reserved内存。差异大于几百MB就要怀疑有没有不对称数据流。4.4 卡间通信速度慢、GPU利用率低下这是多卡训练里最麻烦的一类问题因为它不会给你报错只是训练吞吐量上不去。我排查时习惯按顺序检查这几项DataLoader在处理数据时CPU是不是打满。如果不是说明num_workers太少调大直到CPU利用率接近满分。pin_memoryTrue开了没有。GPU利用率是不是偶尔掉到0然后瞬间打满。如果是大概率是数据预取没跟上优先增大num_workers其次可以尝试prefetch_factor把预取数量调大比如prefetch_factor4。网络通信是不是瓶颈。可以通过torch.profiler分析all-reduce耗时占比如果通信时间占每步总时间比例很高超过30%那就得考虑减少同步频率。不过DDP的bucket已经做了合并一般不是小模型的话不会有太夸张的通信开销。最后检查一下小模型场景。当模型很小时梯度同步的通信时间占比会特别大算力利用率会显得很低。这种情况下如果还是想多卡训练那要考虑把模型的训练改成分组并行或者干脆就单卡训多卡反而更慢。4.5 保存的模型在单卡上加载时报unexpected key这个坑见得太多了。正如前面2.3里说的DDP的state_dict带module.前缀。很多时候你的代码是这样的if use_ddp: torch.save(model.module.state_dict(), save_path) else: torch.save(model.state_dict(), save_path)问题就出在如果你有多卡训练出来的是module.state_dict()之后想用单卡模型加载单卡模型没有module前缀自然加载不上。解决办法是保存前统一做一次key处理state_dict model.state_dict() if use_ddp: state_dict {k.replace(module., ): v for k, v in state_dict.items()} torch.save(state_dict, save_path)这样不管训练时用没用DDP保存出来的都是原始模型的key格式后续加载时就不分单卡多卡了。5. 工程化的几条建议多卡训练这件事看起来是“用更多卡跑同一个模型”实际是把原来单进程内的逻辑改成多进程协作协作的边界越多坑越密集。下面是我在实际项目中沉淀下来的几条经验。第一从第一天起就按多卡写法来写训练脚本。就算你手里只有一张卡也建议用local_rank和DistributedSampler的框架把代码写出来只是把world_size设为1跑单卡。这样做的好处是以后扩卡时基本不用改代码只需要改启动命令。另一个好处是你可以在单卡时就把set_epoch、sampler这些逻辑验证好一旦换多卡出问题的概率会小很多。第二把验证和测试逻辑单独抽象出来。训练过程里跑验证集时多卡环境下建议只在rank0上做避免每张卡都跑一遍带来的浪费和干扰。但要注意如果验证集也要用DDP的Sampler一定要保证验证集总样本数能被卡数整除否则最后rank0的数据和其他卡不一样汇总指标会错乱。最简单的方案是验证时不用分布式采样完全在rank0上跑原始DataLoader。第三日志和checkpoint的写入只交给rank0。你不想看到8张卡同时在控制台刷日志更不想每张卡每隔几步就存一次checkpoint把磁盘写爆。我的做法是if dist.get_rank() 0: print(fEpoch {epoch}, Loss {loss:.4f}) if dist.get_rank() 0: torch.save(state_dict, latest_path)第四监控工具要跟上。多卡训练时用nvidia-smi看实时显存没问题但更要看的就是每张卡的GPU利用率曲线。我习惯用watch -n 0.5 nvidia-smi去盯显存和利用率排查瓶颈时非常有效。如果只用单卡训练GPU利用率掉到20%可能还没什么感知但在多卡同步训练中任何一张卡利用率掉到50%以下全局吞吐就掉到接近单卡的水平这是最直观的“同步惩罚”。最后模型不同阶段用不同训练策略。比如你从头训练一个大模型可以放心用DDP大batch开训但如果是微调一个预训练模型尽量不要改batch size最好保持和预训练时接近的局部batch这样BN统计量不会产生大的偏移效果也稳。这两类场景在代码层面几乎一样但在学习率、BN设置、global batch选择上的取舍完全不同。6. 关于DP和DDP别再用了前面虽然主要讲DDP但有必要简单聊一下DataParallelDP因为它是PyTorch最早提供的多卡方案很多老代码还在用。它的原理是单进程多线程主卡负责汇总梯度并更新参数带来几个先天问题主卡的显存和通信压力远高于其他卡导致7张卡在等1张卡干活Python线程受GIL限制并行效率有限模型参数在每个forward前都要broadcast一次模型越大开销越高。所以DP只有在你只是想临时跑个小实验、不想改代码时才勉强用。正经工程里尤其是模型参数超过千万级别时DDP几乎是唯一合理的选择。DDP在PyTorch演进中已经非常成熟加上torch.compile之后可以把DDP的桶同步和算子融合做得更好让通信开销进一步降低。如果你的模型训练迭代速度上不去试试在DDP基础上开torch.compile有时候能白嫖10%-20%的提速。7. 从单卡到多卡最值得记住的五个关键点写作过程中反复讲的细节比较多最后浓缩成五句话都是实操里真正起决定作用的点batch size的语义在DDP中是按卡计算的全局batch size要乘以卡数学习率要相应缩放。用DistributedSampler之后不要再用DataLoader里的shuffleTrue而且每个epoch必须set_epoch()否则数据顺序一直不变。保存模型时统一去掉module.前缀用state_dict而不是整个模型对象。BN在小batch多卡场景下要主动处理要么每卡的局部batch足够大要么转SyncBatchNorm要么换GroupNorm。做任何分布式操作之前先init_process_group否则所有DDP相关写法都会报“Default process group has not been initialized”。按照这套思路去改代码单卡到多卡基本就是一次平滑迁移。碰到报错不要慌多卡训练的报错大多是环境变量、采样器、state_dict这三个维度的问题排查思路清晰的话基本都能在半小时内定位。最后再分享一个个人习惯我会在训练脚本里加一个--local_world_size参数专门用来在本地小规模调试多卡逻辑。用2张卡跑通流程后再到大集群上开30张卡。这样能把“逻辑错误”和“环境错误”分开排查省掉很多在集群上反复试错的时间。多卡训练本质上是一个工程问题不复杂但细节决定成败稳扎稳打就是最快路径。