聊 PyTorch 分布式模式之前先讲个我观察到的现象十个说要上分布式训练的人有八个其实还没到那个阶段。模型在单卡上能训就是慢一点、显存紧一点于是第一反应就是多搞几张卡跑分布式。结果折腾几天环境、通信、数据加载纷纷出问题最后算下来还不如先做混合精度、梯度累积来优化单卡。分布式不是万金油但当你真的遇到两堵墙——显存墙和算力墙——它又是一套必须掌握的通用技能。这篇文章就围绕 PyTorch 分布式训练来写它解决了什么问题底层进程组和通信怎么运作单卡代码改成多卡到底要动哪几处以及我实际调参和排错时的完整思路。适合两类人一类是模型已经大得单卡装不下想上多卡却不知道从哪下手另一类是已经用 DataParallel 试过但对 DDP、torchrun、NCCL 这些词还是半懂不懂想系统梳理一遍。1. 先想清楚卡之间的墙分布式到底在解决什么问题很多人一听到分布式训练就觉得很高端但 PyTorch 分布式模式的本质非常简单让多个进程协作一起完成同一份训练任务。这里有个容易被忽略的先决条件——你得先明确自己到底是被哪一堵墙卡住了否则后面所有方案都是白搭。1.1 显存墙与算力墙你撞的是哪一堵显存墙是最直观的问题模型在单卡上压根放不下。以十亿参数级别的模型为例FP32 模型权重大约占 4GB 显存Adam 优化器需要保存一阶矩和二阶矩这两项直接翻倍再翻倍又是 8GB反向传播时要保存的梯度大约 4GB。光这三项加起来就超过 16GB还没算激活值和中间特征图常见的 24GB 单卡其实已经非常紧张。如果模型再大一两个数量级单卡必然放不下只能往分布式、混合精度、量化或者模型并行等方向走。算力墙则更常见——模型单卡上能放但训练一轮的时间完全不可接受。比如某个实验单卡要跑两周加一张卡理想情况下能压到一周四张卡可能三天跑完。分布式数据并行在这个场景下几乎是零改造的首选方案因为多张卡各自吃不同的数据梯度汇总之后同步更新模型理论加速比接近线性。判断方法很简单先看显存占用。用torch.cuda.max_memory_allocated()或者显存监控命令看一眼如果训练中显存长期超过 70%说明显存墙压力大如果显存很宽裕主要问题是等不起那就是算力墙。两堵墙对应的分布式策略不一样前者可能需要模型并行后者优先考虑数据并行。1.2 数据并行、模型并行、流水线并行到底该怎么选PyTorch 分布式有三种基本范式很多新手会把它们混在一起。数据并行是最容易理解的每张卡都保留一份完整的模型副本喂进去的数据按 batch 维度切分每张卡计算各自的梯度然后通过集合通信把梯度汇总更新所有副本的参数。打个比方相当于多个厨师拿同一份菜谱各自炒各自的锅最后比对火候取平均之后统一调整。适合模型本身放得下、但算力或吞吐不够的场景。模型并行则是在模型内部做横向切分把模型的不同层放到不同卡上数据按顺序流过这些卡。适合单卡放不下完整模型的场景比如超大 embedding 表或上百层的 Transformer。缺点是层之间的数据搬运是串行的卡与卡之间存在依赖关系利用率容易被打折。流水线并行是在模型并行基础上做优化把模型按层切成若干段每张卡负责一段数据按 micro-batch 流水式地通过各段让所有卡尽量同时工作。它解决的是模型并行下同一时刻只有一张卡在干活的尴尬。实现复杂度最高一般业务场景用不到。我的实际建议是绝大多数业务模型先把数据并行吃透。只有当模型单卡完全放不下才考虑切层。不要一上来就学最复杂的方案分布式训练本身是递进式的跑通 DDP 之后再读 FSDP、流水线等内容会轻松很多。1.3 什么时候不该上分布式这个反直觉的结论值得说小任务真的别硬上。我之前接过一个需求模型是两层 MLP数据量几万条日志上显示每轮训练不到一分钟。这时候配分布式环境、调端口、等 NCCL 初始化花的调试时间比训练本身还长提升完全被通信开销吃掉了。类似的场景单卡训完就行。另一个常见误区是拿还没有跑稳的代码直接上分布式。单卡训练 loss 都不收敛、代码里有内存泄漏或者数据集加载本身有 bug多卡只会把问题放大绝不会自动修正。必须先在单卡把数据 pipeline、模型结构、优化器参数都调到一个稳定状态再考虑横向扩展。如果决定上也建议先走逻辑验证用 2 个进程、极小 batch 跑通流程确认梯度一致、loss 曲线正常然后再升级到全量数据和完整多机环境。这条路径我每次都会走一遍几乎没有例外。2. 分布式模式的地基进程组、通信后端与初始化逻辑PyTorch 分布式跑起来之后所有卡死了、通信超时、结果不一致的问题几乎都能在进程组和通信这块找到根因。这里确实是整个机制的地基搞清楚几个概念比会调 API 更重要。2.1 world_size、rank、local_rank 到底是谁分布式训练启动时会创建一组进程这些进程组成一个通信圈子也叫进程组。在这个组里有几个编号维度必须彻底区分否则配置多机环境的时候很容易懵。world_size 是参与训练的总进程数。如果你有 2 台机器每台 4 张卡那么 world_size 就是 8。rank 是进程全局唯一编号从 0 到 world_size-1。local_rank 则是进程在所在机器内部的编号也相当于当前进程用的是本机哪张卡。还有 node_rank表示当前机器是第几台多机场景下用来区分节点。启动器跟这几者的关系是正常情况下你不会手动去设置这些值torchrun启动时会把它们注入环境变量代码里直接读os.environ[RANK]、os.environ[LOCAL_RANK]就行。老代码里经常出现argparse --local_rank这种手动传参的写法那是torch.distributed.launch时代的遗留风格新项目里最好别学。这里有一个很容易犯错的地方进程的编号和 GPU 的物理编号不一定对应。单机场景local_rank0 用 cuda:0local_rank1 用 cuda:1习惯上没问题但多机或多卡混合配置时如果每台机器的 GPU 编号顺序不统一或者环境变量被重复覆盖就会出现每个进程都在抢同一张卡的诡异现象。所以代码里不能写死cuda:0必须用 local_rank 去映射。2.2 通信后端怎么选NCCL 与 GLOO 的取舍进程组之间的梯度同步和参数广播都要走通信后端。PyTorch 常见的有 NCCL、GLOO、MPI实际上日常项目基本只在前两者之间做选择。后端适用场景特点NCCLGPU 训练Linux 环境性能极高专为多卡集合通信设计all_reduce 等操作做了深度优化GLOOCPU 训练、Windows/macOS 调试易于安装跨平台但 GPU 通信效率远不如 NCCLMPI特定高性能计算集群依赖外部 MPI 环境配置繁琐一般不用单机多卡强烈建议用 NCCL。NCCL 内部会根据卡间拓扑NVLink、PCIe自动选择通信路径多机环境还能利用 RDMA 网络加速。很多人以为单机多卡不需要选后端其实默认走 GLOO 的话小模型看不出差别一旦梯度量变大速度差距会非常明显。CPU 调试时则建议 GLOO。我经常在本地先用 CPU 跑 2 个进程验证分布式逻辑所有数据放到 CPU后端写 gloo跑通之后再切回 GPU。这样做的好处是排错简单——如果 GPU 环境死活起不来至少能确认是代码逻辑问题还是后端配置问题。记住一个基本原则能走 NCCL 就走 NCCL遇到不支持 NCCL 的环境再降级 GLOO。不要因为一时后端报错就随便换NCCL 的报错信息虽然劝退但绝大多数情况是环境问题换后端只是掩盖症状。2.3 init_process_group三种初始化方式背后的共同逻辑所有进程必须互相知道对方的存在才能组成通信圈子。PyTorch 通过torch.distributed.init_process_group完成这个初始化通信方式的配置参数叫 init_method常见三种。env://是最推荐的方式进程从环境变量里读取MASTER_ADDR、MASTER_PORT、WORLD_SIZE、RANK等配置这些变量由torchrun自动注入。代码长得像这样import torch.distributed as dist import os dist.init_process_group( backendnccl, init_methodenv://, )这里 world_size 和 rank 可以不传因为环境变量里已经有了。不过为了明确也可以手动传入。tcp://是手动指定一个节点的地址和端口其他进程连过来。适合自己写启动脚本、不走torchrun的场景。传参方式类似init_methodtcp://10.0.0.1:23456。这种方式的缺点在于一切自己管端口被占用、地址写错都是常见的坑。file://则是通过一个共享文件来同步各进程的信息。所有进程必须能访问同一个文件系统路径适合部分限制网络广播的环境。但它对共享存储有强依赖如果各机器的路径不一致初始化就会失败。我的建议很简单用torchrun启动用env://初始化不要在训练脚本里手写 IP 和端口。把谁启动、谁负责注入配置这件事交给启动器分布式脚本会干净很多。2.4 集合通信原语在训练链路中扮演的角色进程组建好之后进程之间通过集合通信原语交换数据。训练中最常用到几个broadcast 把某个进程的数据广播给所有进程。训练开始前需要保证所有进程从同一份初始权重开始所以 DDP 会在初始化时把 rank 0 的模型状态广播出去。all_reduce 是最典型的梯度同步操作每个进程把本地梯度传出去所有进程执行同一个归约默认是求和最终每个进程都拿到聚合后的结果。DDP 每次反向传播结束后都会对梯度做一次 all_reduce。一句话理解每个进程算出自己的梯度然后大家把梯度加起来求平均保证所有副本的模型更新方向一致。all_gather 和 reduce_scatter 在梯度分片ZeRO和部分高阶并行策略中很关键分别对应收集所有进程的数据和把数据按进程分散归约。FSDP 这类方案大量使用它们。# 手动理解 all_reduce所有进程的 tensor 被累加 dist.all_reduce(tensor, opdist.ReduceOp.SUM) tensor / dist.get_world_size()barrier 则是一个同步屏障所有进程都到齐才继续往下走。它不参与数据交换但调试时非常有用——如果怀疑某个进程卡在某个分支在关键节点插一个 barrier就能定位是哪个进程没到位。3. 从单卡改到多卡的最小改动清单DDP 实操记录术语和原理都清楚了接下来就是动代码。我建议直接学DistributedDataParallel也就是 DDP因为这是 PyTorch 当前最主流的分布式数据并行方案也是你从单卡迁移到多卡路径上成本最低的一种。3.1 DataParallel 为什么不建议再学很多老教程还在讲DataParallel实际上这是一个逐渐被淘汰的方案。DP 是单进程多线程在模型 forward 阶段把数据分成多份、分散到多张卡、再汇总到主卡。这个过程的通信逻辑在主卡上汇合导致主卡负载远高于其他卡显存和算力都不均衡。再加上 DP 的梯度同步用的是 GPU 内部的 reduce在大规模、频繁通信的场景下瓶颈很突出。DDP 则是真正的多进程模型每个 rank 一个独立进程有自己的解释器和模型副本数据由各自进程独立加载。梯度同步通过后端通信直接在所有进程之间完成主卡没有额外的聚合负担。从实际实验看DDP 在同等硬件上的扩展性远好于 DP而且显存分布更均匀。所以我的结论很干脆不需要花时间去调试 DP直接从 DDP 开始。PyTorch 新版已经不建议教学场景继续用 DP 了这个趋势非常明显。3.2 代码改造的六处关键点单卡脚本改成 DDP核心改动就六处。第一初始化进程组。在脚本开头调用dist.init_process_group并把当前进程绑定到对应的 GPU 上。绑定方式是用torch.cuda.set_device(local_rank)替代原来的torch.device(cuda)写死。第二数据加载器必须配合分布式采样器。直接用原来的DataLoader会导致每个进程都吃完整份数据等于是每张卡跑一个独立小训练梯度同步就失去了意义。必须换成DistributedSampler它负责把数据集按世界大小切成互不重叠的份每个 rank 只拿自己那一份数据。第三模型包装成 DDP。把模型移动到local_rank对应的设备后再用DistributedDataParallel包一层之后所有 forward 和 backward 调用方式保持不变。第四随机种子和初始化方式在所有进程间一致。一般情况下 DDP 会自动把 rank 0 的参数同步到其他进程所以不必手动 broadcast。但如果模型里有一些自定义缓冲变量或没有经过 DDP 管理的参数需要在初始化后手动dist.broadcast一次。第五日志打印和 checkpoint 保存只在 rank 0 执行。否则每个进程都会打印一遍日志调试时日志刷屏每个进程都保存模型的话还会互相覆盖。所有只做一次的操作都包一层if not dist.is_initialized() or dist.get_rank() 0:判断。第六每个 epoch 开始时要调用sampler.set_epoch(epoch)。这一步很容易漏。DistributedSampler 内部靠 epoch 数字来决定 shuffle 顺序不设置的话每个 epoch 数据顺序完全相同训练效果会受明显影响。from torch.utils.data import DataLoader, DistributedSampler sampler DistributedSampler(dataset, shuffleTrue) loader DataLoader(dataset, batch_size32, samplersampler, num_workers4) for epoch in range(epochs): sampler.set_epoch(epoch) for batch in loader: ...3.3 用 torchrun 启动而不是自己硬编码DDP 脚本的启动方式直接决定环境变量是否正确。新的标准做法是torchrun它是 Python 内置启动器的升级版自动注入 RANK、LOCAL_RANK、WORLD_SIZE 等环境变量还支持容错重启。单机多卡的最简启动命令torchrun --nproc_per_node2 --nnodes1 train.py--nproc_per_node表示每台机器启动几个进程通常等于可用 GPU 数--nnodes表示节点数单机就是 1。多机场景需要在每个节点分别执行命令并指定主 rendezvous 节点# 主节点rank 0 torchrun --nnodes2 --nproc_per_node8 --rdzv_endpoint192.168.1.10:29500 train.py # 从节点rank 1 torchrun --nnodes2 --nproc_per_node8 --rdzv_endpoint192.168.1.10:29500 --node_rank1 train.py这里--rdzv_endpoint填的是主节点 IP 和端口所有从节点都要指向同一个地址。老教程里用python -m torch.distributed.launch的方式在 PyTorch 新版本里已经标记弃用不建议再复制。还有一个容易被忽略的点代码里不要再写args.rank、args.local_rank这类手动参数。早期版本launch会往脚本注入--local_rank但torchrun不再这样做。如果你在网上找的代码还在用 argparse 读 local_rank通常要改成从os.environ读取。3.4 checkpoint 在分布式下的保存姿势单卡训练保存模型很简单但在 DDP 里有坑。首先是要明确只在 rank 0 保存否则多个进程同时写同一个文件会出现文件互相覆盖甚至写入半截。保存内容用model.module.state_dict()而不是model.state_dict()因为 DDP 包装后的模块多了一层module属性。加载模型时则要看场景。如果是在分布式训练中途恢复 checkpoint所有进程都要先加载同一份权重再进入训练。如果只是用一个训练好的单卡权重去加载 DDP 模型需要处理键名里的module.前缀——要么包装后在对应结构里加载要么把 state_dict 的键名去掉前缀。import torch if dist.get_rank() 0: torch.save(model.module.state_dict(), checkpoint.pt) # 加载时所有进程执行 state_dict torch.load(checkpoint.pt, map_locationfcuda:{local_rank}) model.module.load_state_dict(state_dict)只是有一点要提示DDP 包装内部的参数名会自动带上module.如果模型还包了别的组件键名可能嵌套更多前缀。处理时宁可多打印几行键名核对也别直接硬灌否则容易在训练中途静默地把权重载错。3.5 一份可直接复现的最小骨架代码把以上所有改动合并起来就是下面这段可用于初测的骨架。它不包含具体网络结构只负责把多卡数据加载、模型包装、同步保存这几个流程串起来。import os import torch import torch.distributed as dist import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler def main(): local_rank int(os.environ[LOCAL_RANK]) world_size int(os.environ[WORLD_SIZE]) rank int(os.environ[RANK]) torch.cuda.set_device(local_rank) dist.init_process_group(backendnccl, init_methodenv://) dataset ... # 替换成你的数据集 sampler DistributedSampler(dataset, num_replicasworld_size, rankrank, shuffleTrue) loader DataLoader(dataset, batch_size32, samplersampler, num_workers4, pin_memoryTrue) model nn.Sequential( nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10), ).cuda(local_rank) model DDP(model, device_ids[local_rank]) optimizer torch.optim.SGD(model.parameters(), lr0.1) for epoch in range(5): sampler.set_epoch(epoch) for x, y in loader: x, y x.cuda(local_rank), y.cuda(local_rank) optimizer.zero_grad() loss model(x) loss loss.sum() loss.backward() optimizer.step() if rank 0: torch.save(model.module.state_dict(), fepoch_{epoch}.pt) dist.destroy_process_group() if __name__ __main__: main()注意这个骨架里的 loss 计算是示意性的实际任务要替换成合适的损失函数与标签计算。跑起来之后可以先观察显存占用是否均匀、每个进程的输出是否独立无误后再接真实模型。4. 跑起来之后我踩过的坑以及一条顺畅的排查链路分布式模式跑通只是开始。坦白说真正的知识和经验都在调试阶段。我把自己在这些年调试中踩过频率最高的几类坑整理出来附带完整的排查思路希望能帮你在踩坑时少走弯路。4.1 卡在 TimeoutNCCL 超时的第一根因最典型的现象是训练刚开始一切正常某个 epoch 突然停顿几分钟然后报NCCL timeout。第一次遇到时我非常困惑——代码明明没变为什么越跑越慢后来我才把根因模型梳理清楚。NCCL 超时的本质是某些进程已经进入了通信等待但另外一些进程迟迟没有到场。梯度同步是全局操作只要有一个进程没执行到 all_reduce其他所有进程都会阻塞在通信点。而迟到的原因通常不是 GPU 算力差异而是数据加载不均、磁盘 IO 波动、或者某个进程触发了慢日志。排查链路应该这么走先在init_process_group里显式设置 timeout比如timeoutdatetime.timedelta(minutes10)把无声的挂死变成几分钟后明确报错便于定位。然后在可疑代码段前后插入dist.barrier()看代码卡在哪个位置。如果前一个 barrier 过了、后一个没过说明两个 barrier 之间的逻辑在部分进程上耗时异常。最常见的是数据加载逻辑里有网络请求、磁盘扫描这类不可控 IO。最后检查代码里是否出现条件分支下的通信不匹配。比如某个进程因为提前 break 跳过了 backward其他进程却还在等它的梯度。分布式训练对所有进程执行同样顺序的通信操作要求非常苛刻任何只在一个 rank 内生效的提前返回都会破坏整个协调。4.2 显存不均衡大多不是通信问题是数据分配问题DDP 跑起来之后可能会看到某张卡显存占用明显高于其他卡。很多人第一反应是通信有问题其实大多数情况是数据没有按预期打散。这类问题的排查顺序如下。第一检查DistributedSampler的num_replicas和rank是否和世界大小一致。如果 sampler 还是默认的 sequential 模式每个进程都会加载全部数据等于在资源上做无用功。第二检查 batch size 的处理。分布式数据并行的总 batch size 单卡 batch size × world_size。如果训练脚本里还有全局 batch size 概念你需要想清楚学习率是否要跟着线性缩放。多卡环境下学习率通常按世界大小放大但具体倍率要依据任务调整。第三检查模型是否被绑定到了某个固定 GPU。你可以在模型代码里搜一下有没有to(cuda:0)之类写死的设备字符串。多卡环境下所有设备相关操作都应该由LOCAL_RANK动态推导。第四关注 BatchNorm。DDP 默认各进程独立计算 BN 统计量如果数据分布差异大不同进程的 BN running_mean/running_var 会明显偏离。可以用torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)把普通 BN 转成同步 BN让每个进程在 forward 时交换统计量。代价是增加通信量但收敛稳定性往往更好。4.3 多机通信端口、网卡与存储同步问题单机多卡跑通之后再加一台机器问题会突然变多。很多坑其实与深度学习无关而是分布式系统的基础问题。端口不一致是最常见的主节点的--rdzv_endpoint端口在从节点不可达或者防火墙挡掉了 UDP 通信。排查方法很简单在从节点用 telnet 或类似的工具测试主节点的端口连通性别一上来就怀疑 PyTorch 代码。网卡选择也值得注意。有些机器有多块网卡分布式通信默认走的网卡可能不是高速互联的那块。如果你发现多机训练速度远低于预期先检查是不是通信走了慢速网卡。可以通过环境变量指定使用的接口比如把NCCL_SOCKET_IFNAME设成正确的网卡名。多机训练还依赖共享文件系统。如果每台机器读的是各自本地的数据集副本而数据文件的顺序或切分逻辑不一致模型的收敛会出现诡异差异。最稳妥的方案是所有节点共享同一个数据目录数据集只存一份。另外多机的机器性能和驱动版本最好保持一致。某次我在 4 卡混合异构环境里跑两张卡是新的 Ampere 架构另外两张是旧架构结果无论如何调参训练速度都被两张旧卡拖垮。后来才意识到不同架构的卡根本无法在单 DDP 组里获得均衡性能这个问题只能通过硬件层解决软件优化空间非常有限。4.4 判断分布式训练是否正确的朴素方法分布式调试里最棘手的问题是代码不报错loss 也在降但你不知道它是不是真的按预期工作。我的经验是用三个朴素手段交叉验证。先用极小的数据集和少量进程跑通。数据集只用几十条batch 设很小目标是快速观察整个流程是否正常。这时候不要追求训练效果只看能不能跑完一个 epoch。再和单卡 baseline 对比 loss 曲线。固定随机种子、相同数据顺序单卡跑一份多卡跑一份前几个 step 的 loss 应该基本吻合。如果有明显偏差可能是学习率没有按世界大小缩放或者某个进程的数据没打散。最后直接比对模型权重。训练结束后让所有进程把模型的state_dict打印出前几层参数的 L2 范数理论上应该完全一致。如果不一致基本可以断定梯度同步或初始化广播出了问题需要回到进程组通信层面排查。5. 跑快一点分布式叠加混合精度与数据加载优化DDP 本身解决了能跑的问题接下来的价值在于跑得更快和跑得更稳。分布式训练优化不是单点的事从计算到通信、从加载到存储每一环都可能成为瓶颈。5.1 AMP 与 DDP 叠加要注意什么混合精度训练在现代 GPU 上几乎是标配。PyTorch 的torch.cuda.amp提供自动混合精度接口核心思路是前向和反向计算的部分环节用 FP16 加速同时保留 FP32 主权重来保证数值稳定。跟 DDP 结合时逻辑上没冲突。所有进程执行autocast和GradScaler的方式必须一致而且梯度在反向传播后经过缩放DDP 内部的 all_reduce 会在缩放后的梯度上执行因此不会产生误差。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for x, y in loader: optimizer.zero_grad() with autocast(): loss model(x) loss loss.sum() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()不过要澄清一个误区AMP 主要砍掉的是显存占用和计算开销梯度通信量并不会自动减半。梯度在 DDP 中仍是以 FP32 参与 all_reduce 的所以如果觉得通信是瓶颈不要指望仅仅开启 AMP 就能解决。想降低通信成本得走梯度压缩或梯度分片路线。实际收益方面我的经验是显存占用能下降约 30%~40%在受限于显存或计算吞吐的场景下训练总时长缩短非常明显。但如果本身跑在 CPU 上AMP 没有意义。5.2 数据加载是分布式训练里最容易被低估的瓶颈多卡跑起来之后很多人会惊讶地发现 GPU 利用率始终上不去排除了代码问题之后很大的嫌疑在数据加载。DDP 的每个进程都需要独立喂数据如果DataLoader的瓶颈GPU 只能在大部分时间里空转等待。调节num_workers是最直接的手段。它表示每个进程用几个子进程做数据加载。建议从 CPU 核心数除以进程数往上调但也不要盲目调到很高因为进程间切换和内存占用会抵消收益。我在实际项目里的经验是num_workers4到8通常是甜点区。pin_memoryTrue能显著减少数据从 CPU 到 GPU 的拷贝时间配合非阻塞传输效果更好。还有persistent_workersTrue可以避免每个 epoch 都重新创建 worker减少随机初始化开销。数据预处理如果能离线完成就离线完成。比如图像缩放、解码、tokenize 这些操作放到训练前批量处理好训练时直接读取预处理结果。实时数据增强如果开销太大考虑用 GPU 侧增强替代 CPU 侧逻辑。loader DataLoader( dataset, batch_size32, samplersampler, num_workers8, pin_memoryTrue, persistent_workersTrue, )5.3 单机单卡也能模拟多卡先用 GLOO 验证逻辑调试分布式代码时如果每调一次都要起好几个 GPU 进程效率会非常低。我习惯的做法是在普通 CPU 环境里用 2 个 GLOO 进程先验证逻辑。做法很简单脚本里把设备改成 CPU后端改成 gloo然后同样用torchrun启动两个进程。模型和数据都在 CPU 上跑虽然速度慢但能完整地走一遍 sampler、DDP 包装、checkpoint 保存、集合通信这些逻辑。因为 GLOO 在 CPU 上也能执行集合通信所以很多分布式相关的 bug 能提前暴露。这套方案尤其适合定位是不是通信后端导致的问题。如果 CPU GLOO 下一切正常切到 GPU NCCL 出问题那大概率是 NCCL 环境或网络拓扑配置的问题如果 CPU GLOO 下就报错那就是代码逻辑问题没必要在 GPU 环境里折腾。torchrun --nproc_per_node2 --nnodes1 train_cpu_debug.py脚本里把backend改成gloo设备映射统一用torch.device(cpu)即可。注意有些 GPU 专属算子比如某些 CUDA 实现在 CPU 调试代码里要替换成普通 PyTorch 算子否则会直接报不支持。5.4 再往深处走FSDP 与梯度分片路线如果模型大到单卡放不下甚至 DDP 也力不从心下一步的方向就是梯度分片。PyTorch 里对应的核心工具是 FSDPFully Sharded Data Parallel。它的思想是把模型参数、梯度和优化器状态全部切分到各进程上而不是每个进程都保存一整个副本需要用到某个参数时才做 all_gather用完再释放。FSDP 对显存的削减效果非常明显但通信模式也更复杂。它不是简单的 all_reduce而是 all_gather reduce_scatter 的组合对集群网络的要求更高。我的建议是先跑通 DDP再尝试把一个大模型放进 FSDP观察显存变化和通信开销逐步理解参数切分带来的同步风险。还有流水线并行和序列并行。前者适合超大 transformer 按层切分后者适合超长序列场景。这些方案在工程上各有取舍但学习曲线比 DDP 陡得多。如果你只是想把一个放得下的模型训练得快点完全不需要碰这些等真正走到放不下这一步再按需学习也不迟。我最后想说的是分布式模式本身不解决所有问题它只是把训练从一张卡扩展到多张卡之后重新描述了资源、数据和通信之间的关系。在我自己的实践中最有效的路线永远是先在单卡把模型和数据处理调到稳定状态再上分布式从 2 卡起步跑通后再加更多卡每次改动只动一个变量验证通过再继续。那种配置好环境一次性跑全量的想法听起来很爽实际操作中往往会卡在某个通信异常上花费的时间远超预期。如果你准备开始接触 PyTorch 分布式训练建议先把torchrun、DistributedSampler、DDP、NCCL这几个词练熟再逐步往里加同步 BN、AMP、FSDP。分布式训练有很多坑但绝大多数坑都遵循同一个模式某个进程在错误的时间做了错误的事。只要养成在每个关键节点确认所有进程状态一致的习惯很多问题自然迎刃而解。