刚把手头一个多卡训练的实验收尾,正好有朋友问起PyTorch单GPU和多GPU训练到底差在哪儿、有哪些坑,索性把这次踩过的细节一次性梳理出来。这篇文章不会去教你怎么装CUDA、怎么跑通一个分类网络,而是聚焦在真正工程化之后绕不开的那些对比和细节:数据怎么喂、梯度怎么同步、batch size怎么算、模型怎么存怎么加载,以及在单卡上写得好好的代码,为什么挪到多卡就翻车。
1. 为什么单卡和多卡训练远不止“改几行代码”那么简单
先给结论:把单卡代码改成多卡,表面上看是加几行分布式初始化、包一个DDP,真正的问题是整套训练的“行为”变了。你原来熟悉的batch size、学习率、BN统计、随机种子、评估指标,在多卡环境下的语义都会发生微妙但致命的变化。
单卡训练的逻辑是这样的:一个进程,一块GPU,一份数据,一个模型。每次迭代读一个batch,前向传播算loss,反向传播算梯度,优化器更新参数。整个过程是串行且确定性的,你设了随机种子,基本能复现(虽然有些算子仍然不是完全确定性的)。
多卡训练则引入了“群体协作”的概念。最常用的DataParallel(DP)和DistributedDataParallel(DDP)虽然名字差不多,但哲学完全不同:
- 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_size=64,那么这个64就是每次迭代喂给模型的实际样本数。多卡DDP训练时,如果你仍然设置batch_size=64,那这里的64是指“每张卡的batch size”,全局batch size实际是64 × 卡数。
这个差异直接决定了你的梯度估计质量、收敛速度、最终精度。很多人喜欢说“多卡训练效果应该和单卡一样”,这句话只在“全局batch size和单卡batch size一致”或者“学习率做了相应缩放”的情况下才成立。
具体来说:单卡上batch size=256跑得稳,你搬到4卡DDP,想保持全局256,就得把每卡batch size设为64。如果你把每卡batch size也设成256,那全局就是1024,学习率完全不调的话,模型大概率从一开始就发散。
解决办法是和学习率联动。最常用的经验法则是linear scaling rule:当batch size变为原来的k倍时,学习率也放大k倍。比如单卡batch=256时lr=0.1,4卡全局batch=1024时,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_memory=True打开,减少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。这就是为什么多卡训练时经常出现:单卡batch=32时BN呜呜地收敛,多卡每卡batch=8、全局实际很大,但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与评估指标:在多卡上是否需要特殊处理
常见loss(CrossEntropyLoss、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.benchmark=False和deterministic=True,大概率能复现训练过程。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(f"Epoch {epoch+1}/{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_port="29500"): dist.init_process_group( backend="nccl", init_method="tcp://127.0.0.1:{}".format(master_port), rank=rank, world_size=world_size, ) torch.cuda.set_device(rank)如果你用torchrun或者python -m torch.distributed.launch --nproc_per_node=4 train.py启动,环境变量LOCAL_RANK、RANK、WORLD_SIZE会自动注入,代码可以更简洁:
import os local_rank = int(os.environ["LOCAL_RANK"]) dist.init_process_group(backend="nccl") 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_size=batch_size_per_gpu, sampler=train_sampler, num_workers=4, pin_memory=True, )注意,DistributedSampler内的shuffle参数有默认为True,你不需要在DataLoader里再传shuffle=True。如果你在DataLoader里也传shuffle=True,会直接报错: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_size |
| sampler | 不设置,DataLoader默认shuffle | 必须用DistributedSampler,且每epoch手动set_epoch |
| shuffle参数 | DataLoader(shuffle=True) | DataLoader里不设置,由sampler控制 |
| num_workers | 按需设 | 建议每卡至少4,最好8,否则CPU会成为瓶颈 |
| pin_memory | 建议True | 强烈建议True,MPI场景下能显著减少传输阻塞 |
| device | cuda:0 或 cuda | 通过local_rank动态指定 |
很多教程喜欢在DDP的DataLoader里保留shuffle=True,然后问为什么报错,其实就是这个原因。
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_node=4 --master_port=29500 train.py在代码里通过local_rank = int(os.environ["LOCAL_RANK"])获取设备号。这个方式省心省力,网上各种教程也都默认支持。
4. 多卡训练常见报错与避坑实测
4.1 “AssertionError: DistributedDataParallel is not intended to be used with DataLoader shuffle=True”这类问题
这个报错很好理解: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_memory=True开了没有。- GPU利用率是不是偶尔掉到0然后瞬间打满。如果是,大概率是数据预取没跟上,优先增大
num_workers,其次可以尝试prefetch_factor,把预取数量调大,比如prefetch_factor=4。 - 网络通信是不是瓶颈。可以通过
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(f"Epoch {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,但有必要简单聊一下DataParallel(DP),因为它是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里的shuffle=True,而且每个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张卡。这样能把“逻辑错误”和“环境错误”分开排查,省掉很多在集群上反复试错的时间。多卡训练本质上是一个工程问题,不复杂,但细节决定成败,稳扎稳打就是最快路径。