多卡训练大模型：从DDP到ZeRO，再到张量/流水线/上下文并行

把多个 GPU 拼起来训练大模型，比想象中难，而且难在“通信”和“显存”这两个字上。我练过的模型越大，越觉得分布式训练不是“加卡就提速”，而是要在数据、模型、序列三个维度上做取舍。这篇文章是大模型训练系列的第四篇，主线很明确：从 DDP 这个绝大多数人都会先接触的数据并行方案出发，讲清它为什么能跑通、为什么显存瓶颈还在；接着看 ZeRO 怎么把冗余的优化器状态和梯度分摊掉；然后再进到张量并行、流水线并行，把单层计算和按层切分的并行方式拆开看；最后聊上下文并行，专门解决长序列训练里矩阵算不动、显存放不下的一类问题。读完你应该能形成一张比较清晰的图：什么场景该用什么并行组合，以及为什么一个 70B 模型训练任务通常会同时叠加好几种并行策略，而不是只靠某一种。

1. 并行到底在拆什么：先建立三个坐标轴

1.1 单卡训练的两个瓶颈不是同一个问题

先说显存。以 7B 参数模型为例，如果用 bf16 做前向和反向，参数本身占 14GB，梯度又是 14GB。这还只是暂时性的状态，真正的大头在优化器：AdamW 要维护一份 fp32 的 master 权重，再加一阶动量 m 和二阶动量 v，每一项都是 4 字节，算下来 7B 乘以 4 字节再乘以 3 份，大概是 84GB。哪怕先不看激活值，单卡完整训练一个 step 就需要 14+14+84=112GB 的显存。这就是为什么很多人发现 7B 模型在 40GB 的卡上做全参微调直接 OOM，不是算子写得不好，而是单卡纯跑完整训练状态本来就很勉强。

算力是另一个维度。就算显存够，前向反向一个 step 要跑 10 分钟和跑 1 秒，训练效率完全是两个世界。所以分布式训练实际上同时在两条战线上干活：一是把模型状态（参数、梯度、优化器）分散或者复制到多张卡，二是把计算量分散到多卡并行执行。

1.2 数据的并行、模型的并行、序列的并行

围绕大模型训练，业界常见的并行方式可以归纳成三类：

• 数据并行：每张卡持有一份完整模型，喂不同的数据批次。DDP 就是这个思路，ZeRO 是对它的改进。
• 模型并行：把模型本身切开，每张卡只负责一部分。按切法又分两种：把某一层的矩阵切成多个块，是张量并行；按层的前后顺序切成多段，是流水线并行。
• 序列/上下文并行：针对序列长度这个维度，把长序列切成多段，每张卡负责一段序列上的计算。上下文并行就是这类思路的代表。

这三个坐标轴并不冲突。真正的大规模训练，往往是数据并行、模型并行、序列并行同时叠加。理解每个方案的切分对象和通信频率，比背一堆名词重要得多。

2. DDP：数据并行的“标准答案”和它的隐藏成本

2.1 一次 backward 里到底发生了什么

DDP（Distributed Data Parallel）是 PyTorch 里最常用的分布式训练接口。它的逻辑很直观：假设有 N 张卡，每张卡分别加载一份完整模型权重，然后各自处理一个 local batch。前向互不干扰，反向算出各自的梯度后，执行一次全局 AllReduce，把 N 份梯度求平均，最后每个进程用完全相同的梯度去更新模型。

关键点在于：AllReduce 不只是“等梯度全部算完再做通信”，而是把梯度划分成一个个 bucket。PyTorch DDP 默认以一定的桶大小（bucket_cap_mb，常见 25MB）为单位，反向传播算完一个 bucket 的梯度，就立刻后台发出这个 bucket 的 AllReduce，不用等所有梯度都算完。这个机制叫通信计算重叠。所以 DDP 的性能好不好，很大程度上取决于你能不能把每个 batch 的样本量调大，让计算时间足够覆盖通信时间。

2.2 Ring-AllReduce 为什么比朴素 AllReduce 好

直接在 N 张卡之间做梯度汇总，最笨的方式是每张卡把梯度发给一个主卡，主卡聚合后再广播回去。这个方案在主卡侧带宽压力是 O(N) 的，卡一多就没法用。

实际 DDP 使用的是 Ring-AllReduce，思路是把所有卡的梯度按设备数切成 N 块，卡与卡之间只和相邻节点交换一块数据，绕成一个环。整个过程分成 reduce-scatter 和 all-gather 两段。每一段都有 N-1 次点对点通信，但每张卡每次传输的数据只是总梯度的 1/N。整体算下来，每张卡的有效传输量大约是模型梯度总量的 2 倍，而不是 N 倍。

这就是为什么 DDP 能把梯度同步做得看起来不太耗时：通信量跟卡数没有线性放大，只跟模型大小有关。实际经验里，如果模型是 7B、梯度是 14GB，那么每个 step 的梯度同步总通信量大概在 28GB 左右。在单机 8 卡、NVLink 互联的场景下，这个量级很快就能跑完；一旦跨机走 RDMA 网络，就要看节点间带宽够不够了。

2.3 DDP 的隐藏成本：显存冗余

DDP 最大的隐藏成本不是通信，而是显存冗余。N 张卡各自保存一份完整的参数、梯度、优化器状态，相当于所有状态重复保存 N 份。回到 7B 模型的例子，单卡完整状态大约 112GB，即便分成 64 张卡，每张卡仍然要 112GB，而不是变成 112/64 GB。

所以 DDP 适合那种“单卡还能勉强放下完整模型”的场景。一旦模型大到单卡放不下，哪怕你有 1000 张卡也没有用，因为每张卡都必须有一份完整拷贝。这时候就要上 ZeRO，它的核心工作就是把 DDP 的显存冗余一层层减掉。

3. ZeRO：把优化器状态、梯度、参数三层冗余逐一干掉

3.1 三个阶段的显存账本

ZeRO（Zero Redundancy Optimizer）来自 DeepSpeed，思路并不是去发明一个新的并行维度，而是把 DDP 里重复保存的那 N 份状态，变成每张卡只保存 1/N，用通信换显存。

以 7B 模型、64 卡为例，我粗略算一下单卡显存需求：

• DDP 基线：参数 14GB + 梯度 14GB + 优化器 84GB，约 112GB。
• ZeRO-1：只分片优化器状态，每卡参数 14GB、梯度 14GB、优化器 84/64≈1.3GB，约 29.3GB。
• ZeRO-2：优化器和梯度都分片，每卡参数 14GB、梯度 14/64≈0.2GB、优化器 1.3GB，约 15.5GB。
• ZeRO-3：参数、梯度、优化器全部按层分片，每卡约 1.7GB，这时完整模型状态相当于被切到 64 张卡上。

注意这里没有算激活值。激活显存另算，而且经常比模型状态还猛，通常要靠梯度检查点、小 batch、序列切分来压。

3.2 通信开销的变化：不是免费午餐

ZeRO 用显存换通信，代价要分开看：

• ZeRO-1 和 ZeRO-2 的通信量基本和 DDP 持平。仍然要做梯度汇总，只是把 AllReduce 换成 ReduceScatter，传输总量没有明显增加。
• ZeRO-3 的通信变多，因为每个 step 在反向之前需要从其他卡 Gather 当前层权重，前向要用完整参数，反向算梯度后再 ReduceScatter 梯度。参数、梯度都要动态流动，所以通信量大概是 DDP 的 1.5 到 3 倍，卡数越多、节点间带宽越差，这个开销越明显。

所以 ZeRO-3 虽然显存省得最狠，但对网络很敏感。跨机训练时最好做到节点内带宽够大，或者开启通信重叠和预取参数之类的能力，否则可能会看到 GPU 利用率长期在 30% 以下。

3.3 DeepSpeed 配置里常见的坑

DeepSpeed 的 JSON 配置里，最常用的 ZeRO-2/3 配置大概是这样的：

{ "train_batch_size": 256, "gradient_accumulation_steps": 2, "optimizer": { "type": "AdamW", "params": { "lr": 3e-4 } }, "zero_optimization": { "stage": 3, "reduce_scatter": true, "contiguous_gradients": true, "overlap_comm": true, "stage3_max_live_parameters": 1e9, "stage3_prefetch_bucket_size": 5e7 } }

这里最容易被忽略的是stage3_gather_16bit_weights_on_model_save。ZeRO-3 下权重是分片存在的，存 checkpoint 时如果不先把权重 Gather 回来，保存出来的模型会是残缺的，看起来文件大小不对，加载后损失直接崩。很多人训练到一半保存模型，恢复再训练发现精度不对，排查了半天最后发现是 checkpoint 里没聚合完整权重。

另一个经验是：ZeRO-3 下不要随意直接访问模型的完整参数。因为参数可能并不在本卡上。如果只是想看一下某个张量的数值，最好用 DeepSpeed 提供的上下文管理接口临时 Gather，看完再释放。否则轻则触发大量 AllGather，重则直接死等。

4. 张量并行：把一层拆到多张卡，通信嵌进每一次前向

4.1 Megatron 式拆分的核心：列并行和行并行

当单层矩阵大到一张卡算不动时，就要用张量并行（Tensor Parallelism）。最经典的是 Megatron 的拆法，把 Transformer 里的线性层拆成两类：

• 列并行：权重矩阵按输出维度切，每张卡算其中一块。算完后如果要恢复完整输出，就需要 AllGather。
• 行并行：权重矩阵按输入维度切，每张卡算一个局部部分，最后需要 AllReduce 把各部分结果累加。

Transformer 里的 MLP 通常先是 h 到 4h 的列并行，再是 4h 到 h 的行并行。Attention 里的 QKV 投影常用列并行，输出投影用行并行。于是每一层 Transformer 大概要经历 4 次集合通信。这意味着张量并行是把通信写在了每一层、每一次前向和反向里，通信频率非常高。

4.2 每次前向到底传多少数据

用一个具体例子感受一下数量级。假设 hidden size 是 4096，一个 micro batch 的 token 数是 4096，那么单个张量在 fp16 下的体积是 409640962 字节，约 32MB。一层 Transformer 如果触发 4 次通信，单层就是 128MB 左右。一个 40 层的模型，一次前向光张量并行通信就是几个 GB 量级。这还是在 batch 不算大的情况下。所以张量并行必须跑在高速卡间互联上，比如单机 8 卡的 NVLink。跨节点做张量并行不是绝对不行，但通信时延会明显拖累每一层计算，扩展效率掉得很快。

这里有个实操原则：优先保证 Tensor Parallel 的通信域不超过一个节点，最好就在 8 卡的 NVLink 域内。例如 8 卡节点，就做 TP=8；如果跨节点，宁可把流水线并行放在节点之间，也不要让张量并行跨交换机。

4.3 序列并行：降低张量并行带来的激活冗余

张量并行一个很麻烦的副作用是：为了做 Attention，每张卡经常需要拿到完整的 hidden 状态，这会放大激活显存。Megatron 后来提出了序列并行（Sequence Parallel），把 LayerNorm、Dropout 这类按 token 独立计算的操作按序列维度切分，只在需要完整 hidden 的地方做 AllGather，在 Attention 之后做 ReduceScatter。

这既能把一部分激活显存压下来，又能减少一部分重复通信。现在很多大模型的训练框架都把这个当作张量并行的默认实现，不只是单纯做 ColumnParallel 和 RowParallel。理解这一点对排查显存问题很有用：明明已经切了张量并行，激活还是涨得离谱，十有八九是序列并行没有开启，或者序列长度被重复复制到了每张卡上。

5. 流水线并行：把网络按层切开，吞吐由 micro-batch 和气泡决定

5.1 朴素切分的问题：排队等待

张量并行解决了“单层算不动”，但模型很深时，单张卡要完整加载和计算所有层，仍然吃力。流水线并行（Pipeline Parallelism）的思路是把网络按层顺序切成多段，每张卡只负责其中一段。比如 40 层网络切成 4 段，第 1 到 10 层在卡 0，第 11 到 20 层在卡 1，以此类推。

朴素切分最直接的问题是流水线气泡。前一个 stage 没算完，后面的 stage 只能空等。如果整体只算一个 batch，4 个 stage 的利用率大概只有 1/4。为了让后面尽快动起来，要引入 micro-batch：把一个大 batch 拆成很多个小批量，前一个 stage 算完一个 micro-batch 就立刻传给下一个 stage，形成流水线。

5.2 气泡率怎么算，micro-batch 选多大

气泡是流水线并行最重要的代价。假设有 p 个 stage，m 个 micro-batch，理想情况下整个流水线要经历 m+p-1 个阶段时间，其中只有 m 段时间是让每个 stage 都在有效计算，其余是启动和排空阶段的空闲。近似气泡率可以写成：

bubble ≈ (p-1) / (m+p-1)

举个例子：p=8，m=16 时，气泡率大概是 7/23≈30%。m=64 时降到 7/71≈9.9%。所以流水线并行要尽可能多做 micro-batch，把气泡率压下去。但 micro-batch 太多也有代价，梯度累积步数变多，显存和调度开销都会上来，需要根据显存容量找一个平衡点。

通常我看到有人一开始就把流水线并行分开到几十张卡上，却没有把 micro-batch 调大，结果吞吐极低，还以为是卡间网络问题。其实先算一下气泡率，把 micro-batch 数量调到 stage 数的 4 到 8 倍，情况往往立刻改善。

5.3 GPipe、PipeDream、1F1B 调度怎么选

GPipe 是经典的 fill-drain 模式，先让整个流水线跑完所有 micro-batch 的前向，再统一做反向。实现简单，但内存峰值高，气泡也偏大。PipeDream 尝试把前向和反向交错执行，但可能导致不同 micro-batch 的梯度更新顺序不一致，影响收敛。

实际训练里更常见的是 1F1B 调度，也就是一个 micro-batch 做前向，紧接着一个 micro-batch 做反向，两者交错推进。1F1B 能显著降低显存峰值，也能在前向和反向间更均匀地铺满计算，是当前很多训练框架里流水线并行的默认调度。选调度时，不要只看理论吞吐，还要考虑你实际使用的是否支持 1F1B 以及梯度和参数状态的同步策略。如果用的是比较老的流水线实现，GPipe 反而更省心，但需要准备更大显存。

6. 上下文并行：长序列场景下，把注意力拆到序列维度

6.1 为什么长上下文需要一个新的并行维度

模型参数量没变，但序列变得很长，比如训练 128K、1M 上下文长度的模型。这时问题主要在注意力计算。标准 Attention 的分数矩阵大小是 序列长度 × 序列长度，序列长度一拉长，显存和计算量都会平方级上涨。纯靠数据并行和张量并行，解决不了序列维度上的超线性膨胀。

上下文并行（Context Parallel）的思路是：把一条长序列沿长度方向切成多段，分给多张卡。每张卡只算自己负责的那一段 query，以及对应的局部 K/V 块。不同卡之间需要交换 K/V 块，才能让每个 query 看到完整的上下文。

6.2 环形注意力：通信和计算可以重叠

最常见的实现是环形注意力（Ring Attention）。把序列切成 N 段，每张卡维护一段 query 和一段 K/V。计算时先算本地的注意力分块，然后把 K/V 按环的方向传给下一张卡，环形转一圈，每张卡累计完整查看一遍所有 K/V 块。每转一轮，当前卡把局部注意力结果和 softmax 的局部统计量合并。因为传输的是 K/V 分块，不是重复整个序列，通信量可控；同时 K/V 在环上传输的同时，当前卡还可以继续计算已经到手的块，实现通信计算重叠。

这里有个数值细节：softmax 的全局归一化不能简单相加。每拿到一个新的 K/V 块，都要记录当前的局部最大值，再对已算出的注意力分数做 rescale。环形注意力工程上难的地方就在这里，分块之间如何维护数值稳定，比想象中容易出错。

6.3 上下文并行和序列并行有什么关系

上下文并行和我们前面提的序列并行不是一回事。序列并通常指配合张量并行，把 LayerNorm、Dropout 这类按 token 独立的激活切分；上下文并行则是真正把 Attention 的序列计算拆到设备间，多个设备共同完成同一个序列的注意力计算。

实际训练长上下文模型时，两者经常同时用。注意力部分用上下文并行，前向其他部分用张量并行或数据并行。上下文并行尤其适合那种“单序列特别长、batch 又不大的情况”。如果序列很短，比如只有 2K 或者 4K，强行用上下文并行反而会因为通信开销而没有收益；但当序列到了 64K 以上，单卡放不下 K/V 或者注意力矩阵算不过来时，上下文并行几乎是必须的选择。

7. 组合与选型经验：分布式策略不是越多越好，而是越顺越好

7.1 常规组合关系表

不同并行方式解决的问题不同，放在一张表里看就很清楚：

7.2 一个 72B 模型训练任务的实际组合

假设一个 72B 模型，单卡 80GB，总共 8 节点 64 卡。72B 参数光是 bf16 权重就是 144GB，单卡绝对放不下，因此必须让模型状态在卡间分散。常见的做法是：

• 节点内做 TP=8，利用 NVLink。
• 节点之间用 PP=2 或 PP=4，把网络按层切成 2 到 4 段，避免把张量并行跨交换机。
• 其余维度走 DP，配合 ZeRO-1 或 ZeRO-2 削减冗余。

比如 64 卡可以设 TP=8、PP=4，那么 DP 维度就是 64/(8*4)=2，再用 ZeRO 把优化器状态分片。每个 micro-batch 的激活值通过梯度检查点控制，global batch 用梯度累积去凑。这种“TP+PP+DP+ZeRO”的组合是大模型训练里非常常见的基础框架。

7.3 启动那一刻，NCCL 的报错和拓扑教训

组合越多，分布式启动越容易出问题。我最常被问到的就是 “NCCL timeout”。原因往往不是网络坏了，而是某个进程没有在预期时间完成初始化，或者是节点间网卡名不一致导致通信域选错。

排查时先开NCCL_DEBUG=INFO，再配合nvidia-smi topo -m看看卡之间的互联拓扑。注意容器环境里网络接口选择，必要时设置NCCL_SOCKET_IFNAME指定实际的业务网卡，避免把虚拟网卡和物理网卡混在一起。还有一个容易忽略的点：卡数多时，torchrun的MASTER_ADDR和MASTER_PORT必须对，否则子进程连不上主节点，整个训练会卡在初始化。

另外不要盲目堆并行度。并行方式越多，通信模式越复杂，调度和内存碎片叠加后，真实吞吐未必高于更简单的组合。我的经验是：先用量化评估工具看单卡算力利用率和通信时间占比，再决定要不要加一种并行方式。一个并行方案如果只带来很小的显存节省，却让通信量翻倍，那它就不值得上。

最后再分享一个很实在的排查习惯：每次改并行配置后，先跑一个极小的模型、几个 step，验证 loss 能正常下降，再上正式规模。张量并行切分不对、流水线阶段切错、ZeRO checkpoint 没 Gather 完整，这些坑在几分钟内就能暴露出来。等真跑了几天再发现，浪费的时间和算力足够让人崩溃。分布式训练的路上，谨慎永远不亏。