搞深度学习这几年我最大的感受就是单机单卡训练的日子越来越回不去了。早期跑个 ResNet、VGG一张 1080Ti 也就扛下来了。但到了今天随便一个大语言模型、多模态模型参数量动辄几十亿几百亿你要是还用一张卡硬顶光是显存就不够看。就算模型勉强塞得下训练周期也可能从“按天算”变成“按年算”。PyTorch 分布式训练就是来解决这两个问题的一是显存不够二是速度太慢。但“分布式”这个词听起来挺玄乎网上资料也多是零散片段新手很容易在环境配置、启动方式、模型切分这些环节上反复踩坑。这篇指南我打算用一套完整的思路把策略选型、环境搭建、代码实现、模型选型逻辑串起来讲尽量让一个接触过 PyTorch 但没搞过分布式的读者照着走也能把多卡训练跑起来。内容会覆盖三块核心策略上怎么选DDP、DeepSpeed、模型并行这些方案到底差别在哪实现上怎么改从单卡脚本改造成分布式脚本的关键环节以及模型选型上怎么判断一个模型适不适合用分布式、用哪种分布式最划算。适合的人群也比较广要么你是刚接触分布式的小白想快速搞清楚原理和落地路径要么你已经在单卡上跑通了训练脚本但想迁移到多卡环境提升效率再要么你正在做大模型相关项目需要在 DeepSpeed 和原生 DDP 之间做个取舍。看完这篇至少能让你少走两三天弯路。1. 项目概述为什么分布式训练成了刚需1.1 单卡时代的三大瓶颈先说显存瓶颈。很多人对显存需求没有一个直观概念我习惯用一个粗略的估算公式模型训练总显存约等于参数显存、梯度显存、优化器显存、激活值显存四部分之和。以 7B 参数的大模型为例如果用 FP16 混合精度训练参数本身占大约 14GB7B × 2 Bytes梯度和 FP32 主权重各占 14GB 和 28GBAdam 优化器的动量项和方差项又是 28GB 加 28GB光是这几项常量就已经超过 100GB还没算激活值。单张 A100 80GB 连模型本体都放不下这就是显存瓶颈最直接的体现。再说速度瓶颈。即便一个模型勉强塞进单卡训练速度也常常让人崩溃。我实测过一个中等规模的视觉模型单卡训练一个 epoch 要 40 分钟跑 300 个 epoch 就是 200 小时相当于连续训练 8 天多。这个周期在做研究、调参、跑消融实验时完全不可接受。多卡并行可以把时间压缩到原来的几分之一这对于需要快速迭代的团队来说是刚需中的刚需。最后是资源利用率的问题。很多团队手里其实有好几张卡但平时只跑一个单卡任务剩下卡全在闲置。分布式训练是把闲置算力盘活的最直接方式。不过这里我要多说一句不是所有任务都适合上分布式后面我会专门讲哪些情况不建议用。1.2 分布式到底能解决什么问题分布式训练从目标上分解决的是两类问题一类是“装不下”一类是“跑不动”。“装不下”对应的是模型并行Model Parallelism和流水线并行Pipeline Parallelism。大体思路是把模型切分成多个部分分别放到不同的 GPU 上各个部分协同完成前向和反向计算。这样单卡放不下的大模型可以通过多卡拼接的方式突破显存上限。典型的例子是 GPT-3、LLaMA 这类大语言模型的训练往往同时用到数据并行、模型并行和流水线并行。“跑不动”对应的是数据并行Data Parallelism也就是每张卡上都放一份完整的模型副本各自处理不同的 batch 数据然后通过梯度同步来保持所有副本参数一致。PyTorch 里的 DistributedDataParallelDDP就是这个思路的成熟实现。数据集够大、单卡能放下模型的时候数据并行是最省心、收益最直接的方案。在实际项目中“装不下”和“跑不动”往往同时出现所以现代分布式训练框架都会把多种并行策略组合在一起使用这也是为什么模型选型不能只看参数量还要看整个训练方案怎么编排。1.3 什么情况不建议上分布式这里我要泼一点冷水。分布式不是万金油有些场景硬上分布式反而会降低效率。第一种情况是模型特别小、数据量也不大。我曾经帮朋友调过一个小分类网络单卡训练 20 分钟就完事折腾分布式光环境配置就花了一下午最后多卡同步的通信开销比计算时间还长完全得不偿失。第二种情况是单卡显存充足、训练任务本身是短平快的小实验这类场景保持单卡脚本更灵活。第三种情况是硬件环境不满足基本要求比如多卡之间通过普通千兆网络连接数据同步的带宽会成为严重瓶颈分布式训练的效率可能比单卡还低。判断是否该上分布式我个人的经验是先算一笔账总显存需求是否超过单卡显存、训练总时长是否需要压缩到原来的三分之一以下、多卡之间是否有 NVLink 或高速以太网。三个条件至少满足两个才值得投入精力做分布式改造。2. 方案选型PyTorch分布式训练的几条路线2.1 DP与DDP名字像差距天壤之别PyTorch 很早就有 DataParallelDP这个模块写法极其简单把模型用nn.DataParallel包一层就行。但 DP 在实现上有一个致命问题它是单进程多线程的模型所有卡的数据都要通过主卡上的线程来分发和聚合主卡容易变成性能瓶颈而且受 Python GIL 限制扩展效率非常差。DDPDistributedDataParallel则是多进程方案每个 GPU 由一个独立的 Python 进程负责进程之间通过 NCCL 等后端直接通信不存在 GIL 干扰梯度同步效率远高于 DP。在实际训练中DP 在 4 卡以上的加速比基本就徘徊在 2 到 3 倍左右而 DDP 在 8 卡场景下仍然能保持接近线性的加速效果。所以我给所有朋友的建议都一样新项目直接上 DDPDP 那个口子就别开了。这里把两者的差异整理成一张表方便对照维度DataParallel (DP)DistributedDataParallel (DDP)进程模型单进程多线程多进程每卡一个进程通信方式主卡聚合再分发进程间直接通信NCCL/GlOO受GIL影响明显无扩展性4卡以上衰减严重8卡、多机保持近线性推荐程度不推荐新项目使用主流选择2.2 数据并行、模型并行、流水线并行的取舍数据并行适合单卡能装下模型的场景核心通信发生在每个 step 的梯度同步上。你只要保证数据加载和梯度通信不互相阻塞整体效率就会很稳定。模型并行则是为了突破显存上限把模型切成多段放到不同卡上每段只负责自己的计算。它的缺点是通信量巨大因为切分点上的中间激活值需要在前后卡之间传递而且 GPU 利用率容易因为串行依赖而下降。流水线并行是在模型并行基础上的改良把模型按层切分成多个 stage每个 stage 放在一张卡上同时让不同的 micro-batch 在不同 stage 间流水执行从而提升卡间并行度。它的核心思想类似工厂流水线不同 GPU 同时处理不同阶段的数据减少空闲等待。实现层面虽然 PyTorch 官方提供了torch.distributed.pipeliningtorch 2.x 里逐步完善但大多数业务场景更倾向于直接用 DeepSpeed 这类框架因为框架已经把切分和调度封装好了。还需要提一个词叫张量并行Tensor Parallelism这是大语言模型训练里经常用到的一种模型并行方式把矩阵运算本身切到多卡上做比如把权重矩阵按列切分。它比朴素的模型并行粒度更细、通信更频繁通常只在单机多卡、卡间有高速互联比如 NVLink的环境里使用。如果你用的是公有云租的普通多机集群张量并行要谨慎。2.3 框架生态DDP、DeepSpeed、Horovod怎么选PyTorch 生态里分布式训练框架至少有三个热门选择原生 DDP、微软的 DeepSpeed、以及较早流行的 Horovod。DDP 的优势是零额外依赖和 PyTorch 其他模块天然兼容backward 自动触发梯度同步代码改动量最小。适合模型能塞进单卡、想快速享受多卡加速的场景比如 CV 类任务、中小规模的 Transformer 训练或者作为其他高级框架的底层基础。DeepSpeed 主打大模型场景核心卖点是 ZeRO零冗余优化器能把优化器状态、梯度、参数做分布式切分从而把单卡模型容量扩大好几倍。它还有融合 CUDA kernel、Offload把参数卸载到 CPU/NVMe等能力非常适合训练几十亿甚至上百亿参数的大模型。Horovod 则是框架无关的分布式方案TensorFlow、PyTorch、MXNet 都能接早期在多机训练领域有一批忠实用户。不过随着 PyTorch DDP 的成熟和 DeepSpeed 的流行Horovod 的增量价值已经弱化不少新项目我一般不推荐从零引入它除非你的技术栈里有多框架混用需求。框架核心优势适用场景学习成本原生 DDP代码改动小、稳定性高中小模型、通用多卡训练低DeepSpeedZeRO显存优化、大模型友好LLM、超大模型训练中高Horovod框架无关、多框架统一多框架混用老项目中2.4 模型选型的核心评估维度模型选型不是凭感觉挑一个“大模型”就完事而是要先评估几个关键指标。第一是参数量。参数量直接影响显存需求和通信量。我在前面提过那个估算公式实际使用时可以用一句话版本快速估算FP16 混合精度训练时Adam 优化器 梯度 主权重大概需要参数量 × 16 字节的显存再加上激活值乘个 1.2 到 2 的安全系数。假设你的模型是 1B 参数用 FP16 混合精度固定部分就是约 16GB加上激活值一张 48GB 以上的卡才比较从容这时候优先考虑单卡能装下的数据并行方案如果参数量到 7B 甚至更大直接奔着 DeepSpeed ZeRO 甚至张量并行去。第二是模型结构。Transformer 类模型矩阵运算规整切分相对容易适合各种并行策略而一些带复杂控制流的模型比如动态图、树模型、强化学习里的 actor-critic切分起来会很痛苦这类模型往往更依赖数据并行来提速。第三是训练数据规模。数据并行说到底是拿数据吞吐换训练时间只有数据集足够大多卡并行才有意义。一个图像分类数据集只有一两万张图单卡几个小时就跑完真没必要上分布式。第四是显存总量和通信拓扑。多卡之间如果是 NVLink 连接模型并行和张量并行都能跑得动如果只是万兆以太网跨机器互联通信延迟会很高数据并行反而是最稳的方案。3. 环境搭建与前置准备3.1 硬件规划与网络做分布式训练之前先把硬件底数摸清楚。首先是查看机器上有几张卡、卡间是否通过 NVLink 相连Linux 下用nvidia-smi就能看到 GPU 数量和型号想看 NVLink 拓扑可以用nvidia-smi nvlink -s。单机多卡场景最关键的是卡间通信带宽。NVIDIA 的 NVLink 带宽通常是几百 GB/s 级别而 PCIe 4.0 x16 的单向带宽大约 32GB/s。DDP 每次反向传播都要做一次全规约AllReduce操作通信数据量等于模型参数量乘以梯度字节数。模型参数越大通信量越大对底层互联带宽的要求也越高。如果你手头的多卡机器只是通过 PCIe 交换机互联跑大模型通信开销会非常可观。多机多卡场景网络要求更高。DDP 多机训练时不同机器之间的数据交换要走以太网或者 InfiniBand。千兆以太网因为带宽太小极容易成为瓶颈万兆网可以应付中小规模模型真正的高性能训练集群一般都上 InfiniBand 或 RoCE。对于个人和中小企业如果只有普通千兆局域网建议优先考虑单机多卡不要轻易尝试跨多机训练。3.2 CUDA、PyTorch版本配套关系环境问题占了分布式训练踩坑的一半以上。PyTorch 的每个版本都对 CUDA 版本有明确的配套要求装错了就会出现各种奇怪报错比如CUDA driver version is insufficient或者 import torch 后显示 CUDA 不可用。我推荐一套比较稳的版本组合Python 3.10PyTorch 2.1 或 2.2CUDA 11.8 或 12.1。原因很简单这套组合兼容性测试最充分社区讨论最多遇到问题最容易搜到解决方案。有些新出的组合包比如 Python 3.11 PyTorch 2.6 CUDA 12.4虽然性能可能更好但一些第三方库的预编译轮子还没跟上容易踩依赖冲突。安装方式上推荐用 pip 直接从 PyTorch 官方源安装指定好 CUDA 版本。比如安装 CUDA 12.1 版本的 PyTorch 2.1.2pip install torch2.1.2 torchvision0.16.2 torchaudio0.16.2 --index-url https://download.pytorch.org/whl/cu121注意别用默认 PyPI 源安装 torch因为默认源里是 CPU 版本安装完了 CUDA 全不可用。如果你在 CentOS 这类离线环境下部署建议先在有网环境把 Whl 包下载好再拷贝到目标机器用 pip 离线安装具体就是pip download加上pip install --no-index --find-links两条命令。3.3 Conda环境配置步骤我每次新起项目都用 Conda 创建独立环境避免环境互相污染。这个习惯救过我很多次。创建 Python 3.10 环境并激活conda create -n pytorch_dist python3.10 -y conda activate pytorch_dist安装 PyTorch 及相关库后强烈建议先做一次环境验证确保 CUDA、cuDNN、NCCL 都正常python -c import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())如果输出True且device_count和nvidia-smi看到的卡数一致基本说明单卡环境是通的。这还没完还得验证 NCCL 通信可用import torch.distributed as dist dist.init_process_group(backendnccl, init_methodtcp://127.0.0.1:23456, world_size1, rank0) print(NCCL OK)这个测试能提前把 NCCL 相关的问题暴露出来比直接跑训练再排查要省时间得多。3.4 关于CUDA版本与驱动的一个提醒很多新手容易混淆驱动版本和 CUDA 版本。驱动是显卡固件层面的运行环境CUDA Toolkit 是开发者编译运行程序用的工具包。PyTorch 安装时指定的 CUDA 版本指的是它需要驱动的 CUDA 支持版本并不需要你真的安装完整的 CUDA Toolkit。只要nvidia-smi显示的驱动版本支持你需要的 CUDA 版本右上角会有个最高的 CUDA Version直接用官方预编译包即可。举个例子nvidia-smi如果显示CUDA Version: 12.2那你的驱动至少支持到 CUDA 12.2安装 PyTorch 的 cu121 或 cu118 包都没问题。如果驱动太旧比如只支持 CUDA 11.6那只能装 PyTorch 的 cu118 以下版本否则会报驱动不支持的错。这个规则在分布式训练场景尤其重要因为机器多了之后每台机器的驱动版本五花八门统一版本能省掉大量排查时间。4. 核心实现DDP从零到一4.1 改造训练脚本的关键环节假设你已经有一个单卡训练脚本结构大概是数据加载、模型定义、优化器、训练循环、保存模型这几块。改成 DDP 需要在几个特定位置做插入。第一是初始化分布式环境。用init_process_group初始化通信组单机多卡场景指定nccl后端。这一步必须在创建模型之前完成。第二是获取当前进程对应的 GPU 设备。DDP 是多进程模型每个进程负责一张卡代码里通常这样获取local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank)这里的LOCAL_RANK是torchrun启动时自动注入的环境变量指的是当前进程在当前机器上的编号。比如 4 卡单机训练LOCAL_RANK取值就是 0、1、2、3。第三是用DistributedSampler包装数据集让每个进程训练不同的数据分片。注意这个采样器和普通 shuffle 不同每个 epoch 必须手动调用sampler.set_epoch(epoch)保证每个 epoch 的数据分片顺序都不一样否则所有 epoch 的 shuffle 模式固定不变会影响模型收敛。第四是用DistributedDataParallel包裹模型。模型要先.to(local_rank)到对应 GPU再包进 DDP。第五是保存模型。DDP 模式下每个进程都有完整的模型副本如果每个进程都写一次 checkpoint既浪费磁盘又容易文件冲突。正确做法是只在rank 0的进程上保存。整合起来一个最小的 DDP 单机多卡训练主循环长这样import os import torch import torch.distributed as dist import torch.nn as nn from torch.utils.data import DataLoader, DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP def train(): local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) dist.init_process_group(backendnccl) model nn.Linear(128, 10).to(local_rank) ddp_model DDP(model, device_ids[local_rank]) dataset MyDataset() sampler DistributedSampler(dataset) loader DataLoader(dataset, batch_size32, samplersampler) optimizer torch.optim.Adam(ddp_model.parameters(), lr1e-3) loss_fn nn.CrossEntropyLoss() for epoch in range(10): sampler.set_epoch(epoch) for batch in loader: x, y [t.to(local_rank) for t in batch] optimizer.zero_grad() out ddp_model(x) loss loss_fn(out, y) loss.backward() optimizer.step() if dist.get_rank() 0: torch.save(ddp_model.module.state_dict(), fcheckpoint_{epoch}.pt) dist.destroy_process_group() if __name__ __main__: train()这段代码看起来改动不多但每一步都有讲究。比如DDP(ddp_model, device_ids[local_rank])里的device_ids在多卡场景必须显式指定否则 DDP 会默认使用 cuda:0所有进程都挤到一张卡上直接 OOM。4.2 启动方式torchrun与多机多卡命令PyTorch 官方推荐的启动方式是torchrun它是python -m torch.distributed.run的封装会自动注入LOCAL_RANK、WORLD_SIZE、RANK这些环境变量免去了手动传入的麻烦。单机 4 卡启动命令torchrun --nproc_per_node4 --master_port29500 train.py多机多卡稍微复杂一点。假设两台机器每台 8 卡总进程数 16第一台机器作为 master 节点# 机器0 torchrun --nnodes2 --nproc_per_node8 --node_rank0 --master_addr192.168.1.10 --master_port29500 train.py # 机器1 torchrun --nnodes2 --nproc_per_node8 --node_rank1 --master_addr192.168.1.10 --master_port29500 train.py这里的master_addr必须是 master 节点node_rank0的 IP 地址所有节点必须能互通 TCP 访问而且--master_port要保证不被防火墙拦截。多机场景我强烈建议先 ping 通、再用 telnet 测端口避免训练启动后卡在初始化阶段。4.3 数据加载与BN的细节用 DDP 后一个最常见的隐性坑是DataLoader的num_workers设置。每个 DDP 进程都会独立创建数据加载线程如果卡数多总的 worker 数会暴涨可能把 CPU 内存耗尽。我一般把num_workers设成机器 CPU 核数除以卡数再乘一个 2 的系数这样既能保证数据加载速度又不会把内存撑爆。BN 层在多卡环境下也有讲究。DDP 默认会在反向传播时同步各个进程的 BN 统计量也就是 SyncBN但默认情况下需要模型里的 BN 层被 DDP 识别为“需要梯度同步的层”。如果你的模型用了大 batch 训练单卡 batch 已经足够大比如每卡 64 以上SyncBN 的意义不大反而增加通信开销。数据并行总 batch 很大时我可以接受局部 BN 统计不精确的问题训练效果差别在几个点以内但通信时间能省不少。具体是否开启 SyncBN要看你的总 batch 和卡数而不是盲目开启。4.4 模型保存加载技巧DDP 保存模型的细节前面提了一句“只在 rank 0 保存”这里再展开讲几个容易出问题的点。保存时要取model.module.state_dict()而不是model.state_dict()。因为 DDP 包装后模型的外层多了一个module属性直接存model.state_dict()会把module.前缀带进 key加载时又要做 json 字符串替换非常麻烦。有些朋友图省事保存时直接存整个 DDP 对象这种 checkpoint 体积巨大而且加载时还必须保持 DDP 包装结构一致完全不推荐。加载 checkpoint 时建议在init_process_group之后再读取模型权重并且给每个进程都保留一次读取操作。有些人的写法是只在 rank 0 加载然后广播给其他进程但 DDP 的机制本身就保证所有进程从相同初始状态开始训练所以最简单可靠的做法是每个进程都加载同一份 checkpoint 文件。这里要注意如果加载发生得太早比如在init_process_group之前就 load 模型会有概率出现多个进程看到不同的初始化状态后续梯度同步就会出问题。5. 实操过程与性能调优5.1 从单卡到多卡的完整迁移实例我拿一个典型的图像分类任务来演示完整迁移过程。原始单卡训练脚本里模型是resnet18数据是 CIFAR-10optimizer 是 SGD学习率 0.1batch size 128。单卡训练 50 个 epoch 大约需要 2 小时。换成 4 卡 DDP 之后我做了这样几处调整每卡 batch size 保持 128所以总 batch 变成 512学习率从 0.1 线性放大到 0.4把DataLoader的 shuffle 参数去掉改用DistributedSampler模型用DDP包裹。迁移完成后的训练时间实测在 35 分钟左右加速比约 3.4 倍没有达到理论上的 4 倍这个损失主要来自梯度同步和数据加载竞争属于正常现象。值得特别注意的是学习率缩放。数据并行把总 batch size 扩大到了原来的 n 倍梯度更新频率和统计噪声都变了。如果不调整学习率模型收敛曲线会明显抖动甚至不收敛。PyTorch 官方推荐的线性缩放法则linear scaling rule是new_lr base_lr × num_gpus。比如单卡 0.14 卡就用 0.4。当然这个规则在大 batch 超过一定阈值后就不再完全适用但作为初始起点非常有效。5.2 梯度累积与混合精度梯度累积可以解决两个问题。一是显存不足当单卡 batch size 调小后每个 step 的梯度噪声变大、BN 统计不稳定通过累积若干个 step 的梯度再更新参数可以模拟出接近大 batch 的效果。二是通信效率DDP 默认每个 step 都做一次梯度同步如果每累积 4 步再同步一次通信次数减少为原来的四分之一通信开销明显下降。混合精度AMP则是分布式训练里另一个性价比很高的优化手段。用torch.cuda.amp.autocast把前向和 loss 计算放在 FP16 下执行用GradScaler对梯度做动态缩放防止 FP16 梯度下溢。显存占用能省下 40% 到 50%而且因为 FP16 下的 Tensor Core 计算速率更快训练吞吐也能提升不少。具体代码大致这样from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for batch in loader: optimizer.zero_grad() with autocast(): out model(x) loss loss_fn(out, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度和 DDP 兼容得很好因为梯度同步发生在backward过程中scaler 对梯度缩放不影响进程间通信。实测 4 卡 AMP 的组合训练速度还能在 DDP 基础上再提升 20% 到 30%是现在模型训练的默认配置。5.3 性能监控与瓶颈定位分布式训练跑起来容易但跑得快不快是另一回事。我通常先用nvidia-smi观察每张卡的利用率。如果GPU-Util长时间在 90% 以下说明系统存在瓶颈。常见的瓶颈有三种。第一种是数据加载瓶颈表现为 GPU 利用率忽高忽低训练过程中 GPU 经常处于等待状态这种要看 CPU 内存和磁盘 IO增加num_workers或者用DataLoader的persistent_workersTrue能缓解。第二种是通信瓶颈表现为 GPU 利用率很高但整卡训练时间比单卡没有显著降低这种要检查卡间互联方式和 NCCL 环境变量。第三种是模型同步等待在大模型场景下比较常见表现为各卡 GPU 利用率不均。定位瓶颈我比较推荐用 PyTorch Profiler 或者简单的 Python profiler。前者功能强大但输出信息量大新手容易看晕。我日常调试用的方法是先打印几个关键时间点数据加载耗时、前向耗时、反向耗时、梯度同步耗时。通过日志里这四项的占比就能很清楚地看出瓶颈在哪块。现象可能瓶颈排查方向GPU利用率低于50%数据加载增大num_workers、检查磁盘速度多卡加速比远低于卡数通信开销检查NVLink、NCCL环境变量单卡正常、多卡OOM内存碎片或默认device检查device_ids设置6. 常见问题与排查技巧实录6.1 典型报错与解决方案分布式训练的报错信息往往比较抽象我整理了三个最高频的错误和处理方式。第一个是CUDA out of memory。这个报错在 DDP 里出现最常见的不是显存真的不够而是device_ids设置不对导致多个进程挤在同一张卡上。排查方式很简单训练时另开一个终端跑nvidia-smi看看是不是某张卡显存打满其他卡空闲。如果是检查LOCAL_RANK地图设置确保每张卡只分配给一个进程。如果真的显存不够就把每卡 batch size 调小配合梯度累积补齐。第二个是NCCL timeout。多机训练时经常出现本质是某个节点没能在规定时间内同步完成。常见原因是机器间网络不通、防火墙拦截、或者某个节点启动失败。排查思路是先用telnet master_ip master_port测试端口连通性再确认每个节点的torchrun命令里--master_addr都指向同一个节点。NCCL 的超时时间也可以通过环境变量调大比如export NCCL_TIMEOUT1800给慢节点多留些缓冲。第三个是RuntimeError: Address already in use。原因是上一次训练进程没完全退出端口还被占用。处理办法是ps -ef | grep python找到残留进程先 kill 掉或者干脆换一个--master_port。6.2 多机部署避坑指南多机分布式训练的复杂度比单机高一个量级很多问题发生在启动阶段。第一个大坑是每台机器的环境不一致。Python 版本、PyTorch 版本、CUDA 版本至少这三样必须完全相同。我曾经遇到一台机器装的是 GPU 版 torch另一台装成了 CPU 版训练启动后模型各算各的梯度同步完全错乱。要避免这个问题建议把所有机器的环境用同一个 Conda yml 文件创建并用pip freeze对比三方库版本。第二个大坑是共享文件系统。DDP 训练需要所有进程能访问同一份数据集和 checkpoint 目录。如果各机器数据分散存放比如机器 0 有 train set 前 80%机器 1 有后 80%那实际上两边都在用不同子集训练模型参数无法正确同步。最稳妥的做法是把数据集放在共享存储如 NFS上或者先在每台机器本地同步一份完全相同的副本。第三个大坑是主节点故障。多机训练中master 节点一旦崩溃整个训练作业都会挂掉。因此多机场景我会建议在 master 节点用nohup或tmux启动训练并且保留日志输出方便崩溃后恢复。6.3 检查清单与快速验证我每次跑分布式训练前都会照着下面这个清单快速过一遍nvidia-smi确认每张卡状态正常python -c import torch; print(torch.cuda.device_count())确认 PyTorch 能识别全部 GPU单进程跑一个最小训练循环确认模型和数据没问题用torchrun --nproc_per_node2跑 2 卡最小 DDP 测试确认通信正常检查数据集路径在每台机器都可访问确认 checkpoint 目录有写权限检查--master_port未被占用。如果只是一次快速验证我建议直接用一个极小模型比如两层线性层和少样本数据跑通整个 DDP 流程。验证通过后再切回真实模型和全量数据这样能在最短时间内暴露分布式代码的 bug避免在正式训练中反复重启浪费时间。7. 模型选型实战如何评估一个模型该不该上分布式7.1 从参数量和显存需求反推方案模型选型这件事我习惯用一套判断树来思考而不是凭感觉拍板。第一步估算模型训练所需的峰值显存。用前面提到的公式如果是 FP16 混合精度 Adam 优化器固定显存大约为参数量 × 16 字节如果单卡显存足以覆盖这个数值加上激活值优先考虑纯数据并行DDP。如果单卡显存覆盖不了需要进一步判断超出多少。超出的比例小于 2 倍可以用 DeepSpeed 的 ZeRO Stage 2它把优化器状态切分掉显存压力会明显下降。超出比例超过几倍就得上 ZeRO Stage 3 甚至张量并行、流水线并行组合。第二步评估模型结构。Transformer 类模型因为有规整的 attention 结构张量并行的切分点明确社区实践也成熟这类模型优先考虑 DeepSpeed而一些自定义的 CNN 结构如果切分困难就不要强行模型并行先尝试降低精度、优化结构、或用梯度累积把单卡 batch 降下来换显存。第三步评估训练数据规模和总训练时长。数据量如果不够大分布式带来的收益会被通信开销稀释。我实际见过一个项目模型 2B 参数、数据只有 5GB用 8 卡 DDP 训练总时长只比单卡快了 1.6 倍完全没体现出并行优势。后来切回单卡加更高效的数据加载方式反而更快。7.2 什么场景下应该选DDP以外的方案对于中小模型DDP 是绝对的主流。模型大到单卡确实放不下时就要认真考虑 DeepSpeed 或者更大规模的并行策略。我参与过一个大模型相关项目模型参数量在 7B 左右单张 A100 80GB 勉强能塞下模型但没法训练因为激活值和优化器状态会把显存直接打爆。后来我们用 DeepSpeed ZeRO Stage 3把优化器状态和模型参数分散到 16 张卡上配合梯度 checkpointing才把 7B 模型跑起来。这个项目给我的体会是模型规模到了单卡装不下的程度先上 DeepSpeed别急着手写张量并行。DeepSpeed 把这些复杂调度都封装好了代价是你得理解它的zero_optimization配置参数但踩坑数量比手写实现少一个量级。另外说一句张量并行的适用条件。它通常只用于单机多卡、卡间有 NVLink 高速互联的场景。如果跨机器做张量并行每次矩阵乘法都需要做 AllReduce 这类频繁通信以太网带宽根本扛不住性能会断崖式下跌。所以做多机的大模型训练实际用的是“数据并行 × ZeRO × 流水线并行”的组合而不是简单的张量并行扩展。7.3 模型选型的长期视角除了眼前能不能训得动选型还要看后续迭代和维护成本。DDP 的改动最小代码可读性好后续换人维护也容易上手DeepSpeed 功能强但依赖多、配置复杂遇到版本升级容易出问题手写模型并行虽然灵活但开发量大代码难读只在极特殊场景下才值得。我个人的建议是除非你的模型大到了单卡根本装不下的地步否则优先选 DDP只有在 DDP 支撑不了的时候才考虑升级到 DeepSpeed至于更复杂的多级并行组合最好在团队里有资深分布式开发经验的同事指导下进行。模型选型不是追求最炫的方案而是找到当前团队技术栈和硬件条件下性价比最高的一条路。我在实际项目中养成的一个习惯是每接手一个新任务先花半天时间做环境验证和小规模基准测试实实在在测一下单卡到多卡的加速比与显存占用再决定最终的方案。这半天的时间投入通常会帮我在后续训练周期间省下好几天调参排错的时间。分布式训练没有银弹大量工作都体现在这些基础而琐碎的验证环节上。
PyTorch分布式训练实战：从DDP到DeepSpeed的完整指南
搞深度学习这几年我最大的感受就是单机单卡训练的日子越来越回不去了。早期跑个 ResNet、VGG一张 1080Ti 也就扛下来了。但到了今天随便一个大语言模型、多模态模型参数量动辄几十亿几百亿你要是还用一张卡硬顶光是显存就不够看。就算模型勉强塞得下训练周期也可能从“按天算”变成“按年算”。PyTorch 分布式训练就是来解决这两个问题的一是显存不够二是速度太慢。但“分布式”这个词听起来挺玄乎网上资料也多是零散片段新手很容易在环境配置、启动方式、模型切分这些环节上反复踩坑。这篇指南我打算用一套完整的思路把策略选型、环境搭建、代码实现、模型选型逻辑串起来讲尽量让一个接触过 PyTorch 但没搞过分布式的读者照着走也能把多卡训练跑起来。内容会覆盖三块核心策略上怎么选DDP、DeepSpeed、模型并行这些方案到底差别在哪实现上怎么改从单卡脚本改造成分布式脚本的关键环节以及模型选型上怎么判断一个模型适不适合用分布式、用哪种分布式最划算。适合的人群也比较广要么你是刚接触分布式的小白想快速搞清楚原理和落地路径要么你已经在单卡上跑通了训练脚本但想迁移到多卡环境提升效率再要么你正在做大模型相关项目需要在 DeepSpeed 和原生 DDP 之间做个取舍。看完这篇至少能让你少走两三天弯路。1. 项目概述为什么分布式训练成了刚需1.1 单卡时代的三大瓶颈先说显存瓶颈。很多人对显存需求没有一个直观概念我习惯用一个粗略的估算公式模型训练总显存约等于参数显存、梯度显存、优化器显存、激活值显存四部分之和。以 7B 参数的大模型为例如果用 FP16 混合精度训练参数本身占大约 14GB7B × 2 Bytes梯度和 FP32 主权重各占 14GB 和 28GBAdam 优化器的动量项和方差项又是 28GB 加 28GB光是这几项常量就已经超过 100GB还没算激活值。单张 A100 80GB 连模型本体都放不下这就是显存瓶颈最直接的体现。再说速度瓶颈。即便一个模型勉强塞进单卡训练速度也常常让人崩溃。我实测过一个中等规模的视觉模型单卡训练一个 epoch 要 40 分钟跑 300 个 epoch 就是 200 小时相当于连续训练 8 天多。这个周期在做研究、调参、跑消融实验时完全不可接受。多卡并行可以把时间压缩到原来的几分之一这对于需要快速迭代的团队来说是刚需中的刚需。最后是资源利用率的问题。很多团队手里其实有好几张卡但平时只跑一个单卡任务剩下卡全在闲置。分布式训练是把闲置算力盘活的最直接方式。不过这里我要多说一句不是所有任务都适合上分布式后面我会专门讲哪些情况不建议用。1.2 分布式到底能解决什么问题分布式训练从目标上分解决的是两类问题一类是“装不下”一类是“跑不动”。“装不下”对应的是模型并行Model Parallelism和流水线并行Pipeline Parallelism。大体思路是把模型切分成多个部分分别放到不同的 GPU 上各个部分协同完成前向和反向计算。这样单卡放不下的大模型可以通过多卡拼接的方式突破显存上限。典型的例子是 GPT-3、LLaMA 这类大语言模型的训练往往同时用到数据并行、模型并行和流水线并行。“跑不动”对应的是数据并行Data Parallelism也就是每张卡上都放一份完整的模型副本各自处理不同的 batch 数据然后通过梯度同步来保持所有副本参数一致。PyTorch 里的 DistributedDataParallelDDP就是这个思路的成熟实现。数据集够大、单卡能放下模型的时候数据并行是最省心、收益最直接的方案。在实际项目中“装不下”和“跑不动”往往同时出现所以现代分布式训练框架都会把多种并行策略组合在一起使用这也是为什么模型选型不能只看参数量还要看整个训练方案怎么编排。1.3 什么情况不建议上分布式这里我要泼一点冷水。分布式不是万金油有些场景硬上分布式反而会降低效率。第一种情况是模型特别小、数据量也不大。我曾经帮朋友调过一个小分类网络单卡训练 20 分钟就完事折腾分布式光环境配置就花了一下午最后多卡同步的通信开销比计算时间还长完全得不偿失。第二种情况是单卡显存充足、训练任务本身是短平快的小实验这类场景保持单卡脚本更灵活。第三种情况是硬件环境不满足基本要求比如多卡之间通过普通千兆网络连接数据同步的带宽会成为严重瓶颈分布式训练的效率可能比单卡还低。判断是否该上分布式我个人的经验是先算一笔账总显存需求是否超过单卡显存、训练总时长是否需要压缩到原来的三分之一以下、多卡之间是否有 NVLink 或高速以太网。三个条件至少满足两个才值得投入精力做分布式改造。2. 方案选型PyTorch分布式训练的几条路线2.1 DP与DDP名字像差距天壤之别PyTorch 很早就有 DataParallelDP这个模块写法极其简单把模型用nn.DataParallel包一层就行。但 DP 在实现上有一个致命问题它是单进程多线程的模型所有卡的数据都要通过主卡上的线程来分发和聚合主卡容易变成性能瓶颈而且受 Python GIL 限制扩展效率非常差。DDPDistributedDataParallel则是多进程方案每个 GPU 由一个独立的 Python 进程负责进程之间通过 NCCL 等后端直接通信不存在 GIL 干扰梯度同步效率远高于 DP。在实际训练中DP 在 4 卡以上的加速比基本就徘徊在 2 到 3 倍左右而 DDP 在 8 卡场景下仍然能保持接近线性的加速效果。所以我给所有朋友的建议都一样新项目直接上 DDPDP 那个口子就别开了。这里把两者的差异整理成一张表方便对照维度DataParallel (DP)DistributedDataParallel (DDP)进程模型单进程多线程多进程每卡一个进程通信方式主卡聚合再分发进程间直接通信NCCL/GlOO受GIL影响明显无扩展性4卡以上衰减严重8卡、多机保持近线性推荐程度不推荐新项目使用主流选择2.2 数据并行、模型并行、流水线并行的取舍数据并行适合单卡能装下模型的场景核心通信发生在每个 step 的梯度同步上。你只要保证数据加载和梯度通信不互相阻塞整体效率就会很稳定。模型并行则是为了突破显存上限把模型切成多段放到不同卡上每段只负责自己的计算。它的缺点是通信量巨大因为切分点上的中间激活值需要在前后卡之间传递而且 GPU 利用率容易因为串行依赖而下降。流水线并行是在模型并行基础上的改良把模型按层切分成多个 stage每个 stage 放在一张卡上同时让不同的 micro-batch 在不同 stage 间流水执行从而提升卡间并行度。它的核心思想类似工厂流水线不同 GPU 同时处理不同阶段的数据减少空闲等待。实现层面虽然 PyTorch 官方提供了torch.distributed.pipeliningtorch 2.x 里逐步完善但大多数业务场景更倾向于直接用 DeepSpeed 这类框架因为框架已经把切分和调度封装好了。还需要提一个词叫张量并行Tensor Parallelism这是大语言模型训练里经常用到的一种模型并行方式把矩阵运算本身切到多卡上做比如把权重矩阵按列切分。它比朴素的模型并行粒度更细、通信更频繁通常只在单机多卡、卡间有高速互联比如 NVLink的环境里使用。如果你用的是公有云租的普通多机集群张量并行要谨慎。2.3 框架生态DDP、DeepSpeed、Horovod怎么选PyTorch 生态里分布式训练框架至少有三个热门选择原生 DDP、微软的 DeepSpeed、以及较早流行的 Horovod。DDP 的优势是零额外依赖和 PyTorch 其他模块天然兼容backward 自动触发梯度同步代码改动量最小。适合模型能塞进单卡、想快速享受多卡加速的场景比如 CV 类任务、中小规模的 Transformer 训练或者作为其他高级框架的底层基础。DeepSpeed 主打大模型场景核心卖点是 ZeRO零冗余优化器能把优化器状态、梯度、参数做分布式切分从而把单卡模型容量扩大好几倍。它还有融合 CUDA kernel、Offload把参数卸载到 CPU/NVMe等能力非常适合训练几十亿甚至上百亿参数的大模型。Horovod 则是框架无关的分布式方案TensorFlow、PyTorch、MXNet 都能接早期在多机训练领域有一批忠实用户。不过随着 PyTorch DDP 的成熟和 DeepSpeed 的流行Horovod 的增量价值已经弱化不少新项目我一般不推荐从零引入它除非你的技术栈里有多框架混用需求。框架核心优势适用场景学习成本原生 DDP代码改动小、稳定性高中小模型、通用多卡训练低DeepSpeedZeRO显存优化、大模型友好LLM、超大模型训练中高Horovod框架无关、多框架统一多框架混用老项目中2.4 模型选型的核心评估维度模型选型不是凭感觉挑一个“大模型”就完事而是要先评估几个关键指标。第一是参数量。参数量直接影响显存需求和通信量。我在前面提过那个估算公式实际使用时可以用一句话版本快速估算FP16 混合精度训练时Adam 优化器 梯度 主权重大概需要参数量 × 16 字节的显存再加上激活值乘个 1.2 到 2 的安全系数。假设你的模型是 1B 参数用 FP16 混合精度固定部分就是约 16GB加上激活值一张 48GB 以上的卡才比较从容这时候优先考虑单卡能装下的数据并行方案如果参数量到 7B 甚至更大直接奔着 DeepSpeed ZeRO 甚至张量并行去。第二是模型结构。Transformer 类模型矩阵运算规整切分相对容易适合各种并行策略而一些带复杂控制流的模型比如动态图、树模型、强化学习里的 actor-critic切分起来会很痛苦这类模型往往更依赖数据并行来提速。第三是训练数据规模。数据并行说到底是拿数据吞吐换训练时间只有数据集足够大多卡并行才有意义。一个图像分类数据集只有一两万张图单卡几个小时就跑完真没必要上分布式。第四是显存总量和通信拓扑。多卡之间如果是 NVLink 连接模型并行和张量并行都能跑得动如果只是万兆以太网跨机器互联通信延迟会很高数据并行反而是最稳的方案。3. 环境搭建与前置准备3.1 硬件规划与网络做分布式训练之前先把硬件底数摸清楚。首先是查看机器上有几张卡、卡间是否通过 NVLink 相连Linux 下用nvidia-smi就能看到 GPU 数量和型号想看 NVLink 拓扑可以用nvidia-smi nvlink -s。单机多卡场景最关键的是卡间通信带宽。NVIDIA 的 NVLink 带宽通常是几百 GB/s 级别而 PCIe 4.0 x16 的单向带宽大约 32GB/s。DDP 每次反向传播都要做一次全规约AllReduce操作通信数据量等于模型参数量乘以梯度字节数。模型参数越大通信量越大对底层互联带宽的要求也越高。如果你手头的多卡机器只是通过 PCIe 交换机互联跑大模型通信开销会非常可观。多机多卡场景网络要求更高。DDP 多机训练时不同机器之间的数据交换要走以太网或者 InfiniBand。千兆以太网因为带宽太小极容易成为瓶颈万兆网可以应付中小规模模型真正的高性能训练集群一般都上 InfiniBand 或 RoCE。对于个人和中小企业如果只有普通千兆局域网建议优先考虑单机多卡不要轻易尝试跨多机训练。3.2 CUDA、PyTorch版本配套关系环境问题占了分布式训练踩坑的一半以上。PyTorch 的每个版本都对 CUDA 版本有明确的配套要求装错了就会出现各种奇怪报错比如CUDA driver version is insufficient或者 import torch 后显示 CUDA 不可用。我推荐一套比较稳的版本组合Python 3.10PyTorch 2.1 或 2.2CUDA 11.8 或 12.1。原因很简单这套组合兼容性测试最充分社区讨论最多遇到问题最容易搜到解决方案。有些新出的组合包比如 Python 3.11 PyTorch 2.6 CUDA 12.4虽然性能可能更好但一些第三方库的预编译轮子还没跟上容易踩依赖冲突。安装方式上推荐用 pip 直接从 PyTorch 官方源安装指定好 CUDA 版本。比如安装 CUDA 12.1 版本的 PyTorch 2.1.2pip install torch2.1.2 torchvision0.16.2 torchaudio0.16.2 --index-url https://download.pytorch.org/whl/cu121注意别用默认 PyPI 源安装 torch因为默认源里是 CPU 版本安装完了 CUDA 全不可用。如果你在 CentOS 这类离线环境下部署建议先在有网环境把 Whl 包下载好再拷贝到目标机器用 pip 离线安装具体就是pip download加上pip install --no-index --find-links两条命令。3.3 Conda环境配置步骤我每次新起项目都用 Conda 创建独立环境避免环境互相污染。这个习惯救过我很多次。创建 Python 3.10 环境并激活conda create -n pytorch_dist python3.10 -y conda activate pytorch_dist安装 PyTorch 及相关库后强烈建议先做一次环境验证确保 CUDA、cuDNN、NCCL 都正常python -c import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())如果输出True且device_count和nvidia-smi看到的卡数一致基本说明单卡环境是通的。这还没完还得验证 NCCL 通信可用import torch.distributed as dist dist.init_process_group(backendnccl, init_methodtcp://127.0.0.1:23456, world_size1, rank0) print(NCCL OK)这个测试能提前把 NCCL 相关的问题暴露出来比直接跑训练再排查要省时间得多。3.4 关于CUDA版本与驱动的一个提醒很多新手容易混淆驱动版本和 CUDA 版本。驱动是显卡固件层面的运行环境CUDA Toolkit 是开发者编译运行程序用的工具包。PyTorch 安装时指定的 CUDA 版本指的是它需要驱动的 CUDA 支持版本并不需要你真的安装完整的 CUDA Toolkit。只要nvidia-smi显示的驱动版本支持你需要的 CUDA 版本右上角会有个最高的 CUDA Version直接用官方预编译包即可。举个例子nvidia-smi如果显示CUDA Version: 12.2那你的驱动至少支持到 CUDA 12.2安装 PyTorch 的 cu121 或 cu118 包都没问题。如果驱动太旧比如只支持 CUDA 11.6那只能装 PyTorch 的 cu118 以下版本否则会报驱动不支持的错。这个规则在分布式训练场景尤其重要因为机器多了之后每台机器的驱动版本五花八门统一版本能省掉大量排查时间。4. 核心实现DDP从零到一4.1 改造训练脚本的关键环节假设你已经有一个单卡训练脚本结构大概是数据加载、模型定义、优化器、训练循环、保存模型这几块。改成 DDP 需要在几个特定位置做插入。第一是初始化分布式环境。用init_process_group初始化通信组单机多卡场景指定nccl后端。这一步必须在创建模型之前完成。第二是获取当前进程对应的 GPU 设备。DDP 是多进程模型每个进程负责一张卡代码里通常这样获取local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank)这里的LOCAL_RANK是torchrun启动时自动注入的环境变量指的是当前进程在当前机器上的编号。比如 4 卡单机训练LOCAL_RANK取值就是 0、1、2、3。第三是用DistributedSampler包装数据集让每个进程训练不同的数据分片。注意这个采样器和普通 shuffle 不同每个 epoch 必须手动调用sampler.set_epoch(epoch)保证每个 epoch 的数据分片顺序都不一样否则所有 epoch 的 shuffle 模式固定不变会影响模型收敛。第四是用DistributedDataParallel包裹模型。模型要先.to(local_rank)到对应 GPU再包进 DDP。第五是保存模型。DDP 模式下每个进程都有完整的模型副本如果每个进程都写一次 checkpoint既浪费磁盘又容易文件冲突。正确做法是只在rank 0的进程上保存。整合起来一个最小的 DDP 单机多卡训练主循环长这样import os import torch import torch.distributed as dist import torch.nn as nn from torch.utils.data import DataLoader, DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP def train(): local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) dist.init_process_group(backendnccl) model nn.Linear(128, 10).to(local_rank) ddp_model DDP(model, device_ids[local_rank]) dataset MyDataset() sampler DistributedSampler(dataset) loader DataLoader(dataset, batch_size32, samplersampler) optimizer torch.optim.Adam(ddp_model.parameters(), lr1e-3) loss_fn nn.CrossEntropyLoss() for epoch in range(10): sampler.set_epoch(epoch) for batch in loader: x, y [t.to(local_rank) for t in batch] optimizer.zero_grad() out ddp_model(x) loss loss_fn(out, y) loss.backward() optimizer.step() if dist.get_rank() 0: torch.save(ddp_model.module.state_dict(), fcheckpoint_{epoch}.pt) dist.destroy_process_group() if __name__ __main__: train()这段代码看起来改动不多但每一步都有讲究。比如DDP(ddp_model, device_ids[local_rank])里的device_ids在多卡场景必须显式指定否则 DDP 会默认使用 cuda:0所有进程都挤到一张卡上直接 OOM。4.2 启动方式torchrun与多机多卡命令PyTorch 官方推荐的启动方式是torchrun它是python -m torch.distributed.run的封装会自动注入LOCAL_RANK、WORLD_SIZE、RANK这些环境变量免去了手动传入的麻烦。单机 4 卡启动命令torchrun --nproc_per_node4 --master_port29500 train.py多机多卡稍微复杂一点。假设两台机器每台 8 卡总进程数 16第一台机器作为 master 节点# 机器0 torchrun --nnodes2 --nproc_per_node8 --node_rank0 --master_addr192.168.1.10 --master_port29500 train.py # 机器1 torchrun --nnodes2 --nproc_per_node8 --node_rank1 --master_addr192.168.1.10 --master_port29500 train.py这里的master_addr必须是 master 节点node_rank0的 IP 地址所有节点必须能互通 TCP 访问而且--master_port要保证不被防火墙拦截。多机场景我强烈建议先 ping 通、再用 telnet 测端口避免训练启动后卡在初始化阶段。4.3 数据加载与BN的细节用 DDP 后一个最常见的隐性坑是DataLoader的num_workers设置。每个 DDP 进程都会独立创建数据加载线程如果卡数多总的 worker 数会暴涨可能把 CPU 内存耗尽。我一般把num_workers设成机器 CPU 核数除以卡数再乘一个 2 的系数这样既能保证数据加载速度又不会把内存撑爆。BN 层在多卡环境下也有讲究。DDP 默认会在反向传播时同步各个进程的 BN 统计量也就是 SyncBN但默认情况下需要模型里的 BN 层被 DDP 识别为“需要梯度同步的层”。如果你的模型用了大 batch 训练单卡 batch 已经足够大比如每卡 64 以上SyncBN 的意义不大反而增加通信开销。数据并行总 batch 很大时我可以接受局部 BN 统计不精确的问题训练效果差别在几个点以内但通信时间能省不少。具体是否开启 SyncBN要看你的总 batch 和卡数而不是盲目开启。4.4 模型保存加载技巧DDP 保存模型的细节前面提了一句“只在 rank 0 保存”这里再展开讲几个容易出问题的点。保存时要取model.module.state_dict()而不是model.state_dict()。因为 DDP 包装后模型的外层多了一个module属性直接存model.state_dict()会把module.前缀带进 key加载时又要做 json 字符串替换非常麻烦。有些朋友图省事保存时直接存整个 DDP 对象这种 checkpoint 体积巨大而且加载时还必须保持 DDP 包装结构一致完全不推荐。加载 checkpoint 时建议在init_process_group之后再读取模型权重并且给每个进程都保留一次读取操作。有些人的写法是只在 rank 0 加载然后广播给其他进程但 DDP 的机制本身就保证所有进程从相同初始状态开始训练所以最简单可靠的做法是每个进程都加载同一份 checkpoint 文件。这里要注意如果加载发生得太早比如在init_process_group之前就 load 模型会有概率出现多个进程看到不同的初始化状态后续梯度同步就会出问题。5. 实操过程与性能调优5.1 从单卡到多卡的完整迁移实例我拿一个典型的图像分类任务来演示完整迁移过程。原始单卡训练脚本里模型是resnet18数据是 CIFAR-10optimizer 是 SGD学习率 0.1batch size 128。单卡训练 50 个 epoch 大约需要 2 小时。换成 4 卡 DDP 之后我做了这样几处调整每卡 batch size 保持 128所以总 batch 变成 512学习率从 0.1 线性放大到 0.4把DataLoader的 shuffle 参数去掉改用DistributedSampler模型用DDP包裹。迁移完成后的训练时间实测在 35 分钟左右加速比约 3.4 倍没有达到理论上的 4 倍这个损失主要来自梯度同步和数据加载竞争属于正常现象。值得特别注意的是学习率缩放。数据并行把总 batch size 扩大到了原来的 n 倍梯度更新频率和统计噪声都变了。如果不调整学习率模型收敛曲线会明显抖动甚至不收敛。PyTorch 官方推荐的线性缩放法则linear scaling rule是new_lr base_lr × num_gpus。比如单卡 0.14 卡就用 0.4。当然这个规则在大 batch 超过一定阈值后就不再完全适用但作为初始起点非常有效。5.2 梯度累积与混合精度梯度累积可以解决两个问题。一是显存不足当单卡 batch size 调小后每个 step 的梯度噪声变大、BN 统计不稳定通过累积若干个 step 的梯度再更新参数可以模拟出接近大 batch 的效果。二是通信效率DDP 默认每个 step 都做一次梯度同步如果每累积 4 步再同步一次通信次数减少为原来的四分之一通信开销明显下降。混合精度AMP则是分布式训练里另一个性价比很高的优化手段。用torch.cuda.amp.autocast把前向和 loss 计算放在 FP16 下执行用GradScaler对梯度做动态缩放防止 FP16 梯度下溢。显存占用能省下 40% 到 50%而且因为 FP16 下的 Tensor Core 计算速率更快训练吞吐也能提升不少。具体代码大致这样from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for batch in loader: optimizer.zero_grad() with autocast(): out model(x) loss loss_fn(out, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度和 DDP 兼容得很好因为梯度同步发生在backward过程中scaler 对梯度缩放不影响进程间通信。实测 4 卡 AMP 的组合训练速度还能在 DDP 基础上再提升 20% 到 30%是现在模型训练的默认配置。5.3 性能监控与瓶颈定位分布式训练跑起来容易但跑得快不快是另一回事。我通常先用nvidia-smi观察每张卡的利用率。如果GPU-Util长时间在 90% 以下说明系统存在瓶颈。常见的瓶颈有三种。第一种是数据加载瓶颈表现为 GPU 利用率忽高忽低训练过程中 GPU 经常处于等待状态这种要看 CPU 内存和磁盘 IO增加num_workers或者用DataLoader的persistent_workersTrue能缓解。第二种是通信瓶颈表现为 GPU 利用率很高但整卡训练时间比单卡没有显著降低这种要检查卡间互联方式和 NCCL 环境变量。第三种是模型同步等待在大模型场景下比较常见表现为各卡 GPU 利用率不均。定位瓶颈我比较推荐用 PyTorch Profiler 或者简单的 Python profiler。前者功能强大但输出信息量大新手容易看晕。我日常调试用的方法是先打印几个关键时间点数据加载耗时、前向耗时、反向耗时、梯度同步耗时。通过日志里这四项的占比就能很清楚地看出瓶颈在哪块。现象可能瓶颈排查方向GPU利用率低于50%数据加载增大num_workers、检查磁盘速度多卡加速比远低于卡数通信开销检查NVLink、NCCL环境变量单卡正常、多卡OOM内存碎片或默认device检查device_ids设置6. 常见问题与排查技巧实录6.1 典型报错与解决方案分布式训练的报错信息往往比较抽象我整理了三个最高频的错误和处理方式。第一个是CUDA out of memory。这个报错在 DDP 里出现最常见的不是显存真的不够而是device_ids设置不对导致多个进程挤在同一张卡上。排查方式很简单训练时另开一个终端跑nvidia-smi看看是不是某张卡显存打满其他卡空闲。如果是检查LOCAL_RANK地图设置确保每张卡只分配给一个进程。如果真的显存不够就把每卡 batch size 调小配合梯度累积补齐。第二个是NCCL timeout。多机训练时经常出现本质是某个节点没能在规定时间内同步完成。常见原因是机器间网络不通、防火墙拦截、或者某个节点启动失败。排查思路是先用telnet master_ip master_port测试端口连通性再确认每个节点的torchrun命令里--master_addr都指向同一个节点。NCCL 的超时时间也可以通过环境变量调大比如export NCCL_TIMEOUT1800给慢节点多留些缓冲。第三个是RuntimeError: Address already in use。原因是上一次训练进程没完全退出端口还被占用。处理办法是ps -ef | grep python找到残留进程先 kill 掉或者干脆换一个--master_port。6.2 多机部署避坑指南多机分布式训练的复杂度比单机高一个量级很多问题发生在启动阶段。第一个大坑是每台机器的环境不一致。Python 版本、PyTorch 版本、CUDA 版本至少这三样必须完全相同。我曾经遇到一台机器装的是 GPU 版 torch另一台装成了 CPU 版训练启动后模型各算各的梯度同步完全错乱。要避免这个问题建议把所有机器的环境用同一个 Conda yml 文件创建并用pip freeze对比三方库版本。第二个大坑是共享文件系统。DDP 训练需要所有进程能访问同一份数据集和 checkpoint 目录。如果各机器数据分散存放比如机器 0 有 train set 前 80%机器 1 有后 80%那实际上两边都在用不同子集训练模型参数无法正确同步。最稳妥的做法是把数据集放在共享存储如 NFS上或者先在每台机器本地同步一份完全相同的副本。第三个大坑是主节点故障。多机训练中master 节点一旦崩溃整个训练作业都会挂掉。因此多机场景我会建议在 master 节点用nohup或tmux启动训练并且保留日志输出方便崩溃后恢复。6.3 检查清单与快速验证我每次跑分布式训练前都会照着下面这个清单快速过一遍nvidia-smi确认每张卡状态正常python -c import torch; print(torch.cuda.device_count())确认 PyTorch 能识别全部 GPU单进程跑一个最小训练循环确认模型和数据没问题用torchrun --nproc_per_node2跑 2 卡最小 DDP 测试确认通信正常检查数据集路径在每台机器都可访问确认 checkpoint 目录有写权限检查--master_port未被占用。如果只是一次快速验证我建议直接用一个极小模型比如两层线性层和少样本数据跑通整个 DDP 流程。验证通过后再切回真实模型和全量数据这样能在最短时间内暴露分布式代码的 bug避免在正式训练中反复重启浪费时间。7. 模型选型实战如何评估一个模型该不该上分布式7.1 从参数量和显存需求反推方案模型选型这件事我习惯用一套判断树来思考而不是凭感觉拍板。第一步估算模型训练所需的峰值显存。用前面提到的公式如果是 FP16 混合精度 Adam 优化器固定显存大约为参数量 × 16 字节如果单卡显存足以覆盖这个数值加上激活值优先考虑纯数据并行DDP。如果单卡显存覆盖不了需要进一步判断超出多少。超出的比例小于 2 倍可以用 DeepSpeed 的 ZeRO Stage 2它把优化器状态切分掉显存压力会明显下降。超出比例超过几倍就得上 ZeRO Stage 3 甚至张量并行、流水线并行组合。第二步评估模型结构。Transformer 类模型因为有规整的 attention 结构张量并行的切分点明确社区实践也成熟这类模型优先考虑 DeepSpeed而一些自定义的 CNN 结构如果切分困难就不要强行模型并行先尝试降低精度、优化结构、或用梯度累积把单卡 batch 降下来换显存。第三步评估训练数据规模和总训练时长。数据量如果不够大分布式带来的收益会被通信开销稀释。我实际见过一个项目模型 2B 参数、数据只有 5GB用 8 卡 DDP 训练总时长只比单卡快了 1.6 倍完全没体现出并行优势。后来切回单卡加更高效的数据加载方式反而更快。7.2 什么场景下应该选DDP以外的方案对于中小模型DDP 是绝对的主流。模型大到单卡确实放不下时就要认真考虑 DeepSpeed 或者更大规模的并行策略。我参与过一个大模型相关项目模型参数量在 7B 左右单张 A100 80GB 勉强能塞下模型但没法训练因为激活值和优化器状态会把显存直接打爆。后来我们用 DeepSpeed ZeRO Stage 3把优化器状态和模型参数分散到 16 张卡上配合梯度 checkpointing才把 7B 模型跑起来。这个项目给我的体会是模型规模到了单卡装不下的程度先上 DeepSpeed别急着手写张量并行。DeepSpeed 把这些复杂调度都封装好了代价是你得理解它的zero_optimization配置参数但踩坑数量比手写实现少一个量级。另外说一句张量并行的适用条件。它通常只用于单机多卡、卡间有 NVLink 高速互联的场景。如果跨机器做张量并行每次矩阵乘法都需要做 AllReduce 这类频繁通信以太网带宽根本扛不住性能会断崖式下跌。所以做多机的大模型训练实际用的是“数据并行 × ZeRO × 流水线并行”的组合而不是简单的张量并行扩展。7.3 模型选型的长期视角除了眼前能不能训得动选型还要看后续迭代和维护成本。DDP 的改动最小代码可读性好后续换人维护也容易上手DeepSpeed 功能强但依赖多、配置复杂遇到版本升级容易出问题手写模型并行虽然灵活但开发量大代码难读只在极特殊场景下才值得。我个人的建议是除非你的模型大到了单卡根本装不下的地步否则优先选 DDP只有在 DDP 支撑不了的时候才考虑升级到 DeepSpeed至于更复杂的多级并行组合最好在团队里有资深分布式开发经验的同事指导下进行。模型选型不是追求最炫的方案而是找到当前团队技术栈和硬件条件下性价比最高的一条路。我在实际项目中养成的一个习惯是每接手一个新任务先花半天时间做环境验证和小规模基准测试实实在在测一下单卡到多卡的加速比与显存占用再决定最终的方案。这半天的时间投入通常会帮我在后续训练周期间省下好几天调参排错的时间。分布式训练没有银弹大量工作都体现在这些基础而琐碎的验证环节上。