分布式训练：几百张 GPU 怎么协作？数据并行、模型并行、流水线并行

我第一次碰到训练 OOM，是在一台 8 卡 A100 的机器上跑一个 13B 参数的模型。我心想，这玩意儿不是号称能装下吗？结果一个 batch 就报错。后来我才明白，“能装下”指的是模型参数，没算上优化器状态、梯度和中间激活。这个认知差让我老老实实啃了一遍分布式训练。你现在看到的这篇文章，就是我当时希望有人能写给我的。

训练大模型本质上要解决两个问题：一张卡放不下，和一张卡算太慢。放不下，是因为模型参数、梯度、优化器状态（比如 Adam 的 m 和 v）加起来远超显存；算太慢，是因为哪怕能放下，数据量一大，迭代周期长得无法忍受。解决问题的思路很朴素：既然一张卡不够，我们就用多张卡一起干。但怎么“一起干”？这就是数据并行、模型并行、流水线并行的核心区别。

数据并行：每人背一份完整模型，各算各的数据

数据并行的逻辑最简单：每个 GPU 都有一份完整的模型副本，只是喂的数据不同。你把一个 batch 拆成 4 份，分给 4 张卡，每张卡独立做前向传播、计算 loss、反向传播，得到各自的梯度。然后所有卡交换梯度，取平均，再用平均后的梯度更新模型参数。这样，每张卡上的模型副本始终保持一致。

这个“交换梯度取平均”的操作，就是分布式训练里最核心的通信原语：AllReduce。早期实现用参数服务器（Parameter Server），所有 worker 把梯度推给一个中心节点，中心节点求平均再发回来。但中心节点会成为瓶颈，而且带宽压力大。后来 Ring AllReduce 成了主流：GPU 排成一个逻辑环，每个卡只跟左右邻居通信，把通信量分散到所有链路，带宽利用率极高。PyTorch 的 DistributedDataParallel (DDP) 底层就用了 NCCL 实现的 Ring AllReduce。PyTorch DDP 文档

数据并行最直观，但有个致命短板：它要求每张卡都能装下完整模型。对于 13B 以上的模型，即使半精度（FP16）参数就 26GB，加上梯度、优化器状态（Adam 需要额外 2 倍参数大小的 m 和 v），一张 80GB A100 也塞不下。所以数据并行适用的场景是：模型不算太大，但数据量巨大，你想加速训练。一旦模型大到单卡放不下，数据并行就失效了。

模型并行（张量并行）：把一层拆开，放到多张卡上

既然整层放不下，那就把一层切开。模型并行，更准确地说张量并行（Tensor Parallelism），就是把一个层内的参数矩阵切分到多个 GPU 上，每个 GPU 只存一部分参数，但仍能完成完整的矩阵运算。

以 Transformer 中最常见的全连接层为例：Y = W * X，其中 W 是一个大矩阵。切开有两种方式，Megatron-LM 论文里讲得很清楚Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism：

• 列并行：把 W 按列切成 N 份，每个 GPU 存一份 W_i，输入 X 需要广播到所有 GPU，各自计算 Y_i = W_i * X，最后把 Y_i 拼起来（AllGather）。
• 行并行：把 W 按行切，输入 X 也按列切（对应 W 的行切分），每个 GPU 存 W_i 和 X_i，计算 Y_i = W_i * X_i，然后把所有 Y_i 加起来（AllReduce）。

张量并行的好处是，它把超大的矩阵乘拆成了多个小矩阵乘，分别在不同 GPU 上并行计算，显存占用直接减半（或按卡数线性减少）。但代价是每次前向和反向都需要 GPU 间通信，而且这类通信通常很密集。比如列并行最后的 AllGather，需要把所有卡的结果拼起来，通信量跟输出大小成正比。因此张量并行通常只在同一个节点内的卡间做（NVLINK 带宽高），跨节点通信开销太大。

一个常见的误解：很多人以为模型并行就是按层切，其实那是流水线并行。模型并行（张量并行）是层内切分，粒度更细，通信量也更大。

流水线并行：按层切分，像工厂流水线一样

流水线并行的思路是把模型按层切分，GPU 0 负责第 1-10 层，GPU 1 负责第 11-20 层，以此类推。数据像流水线一样流过各个 GPU：GPU 0 算完前 10 层，把中间激活传给 GPU 1，GPU 1 接着算，直到最后一张卡输出 loss。反向传播则按相反顺序回流梯度。

这个想法很自然，但有一个巨大的坑：流水线气泡（Bubble）。想象一下，你把一个 batch 的数据喂给 GPU 0，它忙活一阵，GPU 1 必须等 GPU 0 算完才能开始，所以 GPU 1 有一段时间是空闲的。一个 batch 跑完，第 2 个 batch 又得等 GPU 0 先算。流水线越长，气泡越大，GPU 利用率越低。

GPipe 论文GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism 提出了一个巧妙的解法：引入 micro-batch。把一个 batch 切成更小的 micro-batch，比如切成 4 份。GPU 0 算完 micro-batch 1 后立刻传给 GPU 1，然后马上算 micro-batch 2，而不是等整个 batch 算完。这样 GPU 1 就能更早地开始工作，气泡被挤压到两端。理论上，micro-batch 数量越多，气泡占比越小，但数量又不能无限大，因为每个 micro-batch 都需要保存中间激活，显存会爆。GPipe 同时用活化重计算（activation recomputation）来换显存，在反向时重新计算中间激活而不是存着。

另一种更激进的调度是 PipeDream 的 1F1B（one forward one backward） 策略PipeDream: Fast and Efficient Pipeline Parallel DNN Training：每个 GPU 做完一个 micro-batch 的前向，立刻做另一个 micro-batch 的反向，交错执行，让前向和反向重叠，进一步减小气泡。这个方法实现复杂，但效率更高，后来被许多框架采用。

流水线并行相比张量并行，通信量小得多：只需要在切分层之间传递激活值和梯度，不需要在每个矩阵乘时通信。所以它适合跨节点扩展，通信开销低。但它的调度复杂，有气泡，而且要求每层计算量尽量均衡，否则会出现某个 GPU 长期等待。

一张表说清三种并行

这张表很重要。你可能会问：这三种并行是互斥的吗？当然不是。实际训练大模型时，几乎都是混合并行：节点内用张量并行（高速 NVLINK），节点间用流水线并行（网络带宽有限），最外层再套数据并行（多个节点组复制整个模型）。这就是 Megatron-LM 提出的 PTD-P（Pipeline, Tensor, Data Parallelism）Efficient Large-Scale Language Model Training on GPU Clusters。比如训练一个 175B 的 GPT-3，可能用 8 路张量并行，4 路流水线并行，然后数据并行扩展到 64 个这样的节点组，总共 512 张卡。

我踩过的两个坑，以及它们怎么帮我理解本质

坑一：以为 ZeRO 是模型并行，其实它是数据并行的进化。DeepSpeed ZeRO 论文ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 把优化器状态、梯度、参数分别分片（partition）到不同 GPU，每个 GPU 只存一部分，需要时再通过通信取回。这看起来像模型并行，但通信模式是 AllGather 参数再计算，然后 AllReduce 梯度，本质上还是数据并行的逻辑，只是内存布局优化了。ZeRO-3 甚至能让单卡塞不下完整参数，但计算时仍然需要完整参数，所以它需要大量通信，但巧妙地把显存压力转嫁到了通信上。理解这一点后，我才明白为什么 ZeRO 适合在数据并行框架中叠加，而张量并行和流水线并行是独立的维度。

坑二：通信开销的直觉往往是错的。最早我以为流水线并行通信最少，所以效率最高。实际上，流水线并行的气泡问题如果不处理好，GPU 利用率可能只有 50%。而张量并行虽然通信频繁，但利用 NVLINK 的高带宽在节点内几乎无感。数据并行看似通信量大（梯度同步），但通过分桶（bucket）和通信计算重叠（overlap），DDP 能把梯度同步藏在前向反向计算后面，几乎不额外占用时间。所以效率评估不能只看通信量，要看实际占用的墙壁时间。

几个你可能会问的问题

数据并行只能用于小模型吗？

不完全是。配合 ZeRO，数据并行可以训练千亿参数模型。但纯数据并行（无 ZeRO）确实要求每卡装下完整模型，所以对于百亿参数以上模型，单卡放不下，必须结合模型并行或流水线并行。

模型并行和张量并行是一回事吗？

在学术上，模型并行通常指张量并行，即层内切分。但工业界有时把流水线并行也叫模型并行，容易混淆。建议明确区分：张量并行（层内切分）和流水线并行（层间切分）。

流水线并行的气泡怎么消除？

完全消除不可能，只能减小。GPipe 用 micro-batch 把气泡挤压到计算开始和结束，当 micro-batch 数量远大于流水线深度时，气泡占比趋近于 0。PipeDream 用 1F1B 交错调度，让前向和反向重叠，进一步压缩气泡。但调度的实现复杂度增加，且需要处理好显存。

训练 GPT-4 这种超大模型，主要用哪种并行？

实际上是组合拳。典型的大型训练集群中，节点内通常用张量并行（8 卡 NVLink），节点间用流水线并行（跨机），数据并行则用于扩增更多节点组。此外还会结合序列并行（如 Megatron-LM 的序列并行，把长序列切分到不同 GPU 以减少激活内存），以及 ZeRO 优化。所以没有单一答案，而是混合与调优。

怎么判断我的模型该用什么并行？

一个简单决策树：
1. 如果能单卡放下，直接用数据并行（DDP）。
2. 如果单卡放不下，但单层内参数矩阵很大（比如隐藏维度 12288 以上），优先在节点内用张量并行，节点间用流水线并行。
3. 如果模型层数极多，但单层不大，优先流水线并行。
4. 总显存还是不够？叠加 ZeRO，或者用张量并行度更高。

最后，如果你想动手试，推荐从 PyTorch 的 FSDP（Fully Sharded Data Parallel，类似 ZeRO-3）开始，它把模型参数、梯度、优化器状态都分片了，使用起来和 DDP 一样简单，但能让你在有限显存下跑更大的模型。然后可以尝试 Megatron-LM 的代码，感受一下流水线调度和通信算子。痛苦是暂时的，通透是永久的。