单卡放不下？那就把模型拆开
我最早以为，大模型训练就是多买几张显卡，然后数据并行一跑就完事了。直到我亲自尝试在8张A100上加载一个70B的模型——显存直接爆了，OOM。那一刻我才意识到，大模型训练的真正挑战不是算力，而是‘怎么把模型塞进显存，同时还能高效地算’。
后来我发现，业界为了解决这个问题，搞出了一套组合拳：数据并行 + 张量并行 + 流水线并行，合称 3D 并行。听起来很炫，但真正上手时，你会发现这三个‘并行’放在一起，并不是简单的加法，而是一场通信、内存和计算之间的精密博弈。
数据并行：最简单，但也是最不够的
数据并行是最直观的并行方式：每张卡上存一份完整的模型，数据分成多份，每张卡算自己的那份，然后同步梯度。但问题来了——当模型参数大到单卡显存放不下时，数据并行就没用了。你连模型都加载不了，还谈什么训练？这时候，你必须把模型本身拆开，这就是模型并行。
张量并行：把矩阵乘法切碎
张量并行是模型并行的一种，它把单个层内的矩阵运算切分到多张卡上。比如Transformer的注意力层，Q、K、V矩阵可以按列或按行切分，每张卡算一部分，然后通过 all-reduce 通信合并结果。这种方式的好处是通信量相对可控，但坏处是每次前向和反向传播都要做通信，对卡间带宽要求极高，通常需要 NVLink 或 NVSwitch 这样的高速互联，跨机器的话通信开销会吃掉所有收益。
我最早玩张量并行时，天真地把一个层切分到两台机器上，结果训练速度比单卡还慢。后来才明白，张量并行必须放在同一台机器内，甚至同一块 NVSwitch 下的 GPU 之间，不然通信延迟会让你怀疑人生。
流水线并行：像工厂流水线一样处理层
流水线并行是另一种模型并行，它把模型按层切分，比如前10层放在 GPU0，中间10层放在 GPU1，最后10层放在 GPU2。数据像流水线一样依次通过各 GPU。但这样会产生‘气泡’——GPU0 算完第一批数据，要把结果传给 GPU1，这时 GPU1 可能还在等上一批数据，造成空闲。为了减少气泡，引入了微批次（micro-batch）的概念，把一个大批次拆成多个小批次，让流水线填得更满。
为什么需要 3D 并行？因为单种并行根本不够
单用张量并行，受限于单机内的 GPU 数量（比如 8 卡），没法扩展到更大模型；单用流水线并行，通信开销和气泡限制也很大。于是，Megatron-LM 提出把张量并行和流水线并行结合起来，再叠加数据并行，形成 3D 并行。这样，你可以在一个机器内用张量并行切分大层，跨机器用流水线并行切分层，再跨多个这样的流水线组进行数据并行。这样就把模型切分到了几百甚至几千张卡上。
3D 并行的真正挑战：通信开销与流水线气泡
但 3D 并行带来的是通信拓扑的爆炸式复杂化。张量并行需要高频的 all-reduce，流水线并行需要点对点的发送和接收，数据并行需要在所有数据并行组之间做全局梯度同步。这三类通信如果混在一起，且网络拓扑不匹配，就会导致拥塞和延迟。所以，真正部署 3D 并行，你需要仔细规划每一层并行对应的通信组，让不同的通信流使用不同的物理链路，比如张量并行用 NVLink，流水线用 InfiniBand，数据并行也可以用 InfiniBand 但需要隔离。
很长时间里，我以为 3D 并行就是三种并行的简单叠加，后来读了 Megatron-LM 的论文和代码，才发现，他们精心设计了并行策略的映射，比如张量并行大小、流水线并行大小和数据并行大小需要根据模型架构和硬件拓扑来搜索最优组合。没有银弹，只有一堆超参数要调。
我曾以为的误解：ZeRO 和 3D 并行是一回事吗？
很多人会把 ZeRO 和 3D 并行混为一谈。ZeRO（Zero Redundancy Optimizer） 是 DeepSpeed 提出的，它本质上还是数据并行，但通过分片优化器状态、梯度和参数，让每张卡只存一部分，减少显存冗余。ZeRO 可以单独使用，也可以和 3D 并行结合。比如，你可以用 ZeRO-1 只分片优化器状态，同时用张量并行和流水线并行。ZeRO-3 分片模型参数，就相当于一种模型并行，但通信模式不同。所以，3D 并行和 ZeRO 不是互斥的，它们是从不同维度解决显存问题。
一张表说清三种并行
| 并行类型 | 切分维度 | 通信模式 | 硬件要求 | 适用场景 |
|---|---|---|---|---|
| 数据并行 | 数据批次 | 梯度 all-reduce（每步一次） | 任意网络，带宽要求较低 | 模型能放进单卡时 |
| 张量并行 | 层内矩阵 | 前向/反向 all-reduce（每层多次） | 需要极高带宽，如 NVLink | 单机内扩展大层 |
| 流水线并行 | 模型层 | 点对点发送/接收（微批次边界） | 跨节点网络，带宽要求中等 | 跨机器扩展模型深度 |
关于 3D 并行，你可能还想问的
问：张量并行和流水线并行哪个更好？
答：看模型大小和硬件拓扑。张量并行适合单机多卡，流水线并行适合跨机。实际中，你会两者都用。
问：3D 并行能让训练速度线性加速吗？
答：几乎不可能。通信开销、气泡和负载不均衡会让实际加速比远低于卡数。比如 1024 卡训练，可能只有 500 卡的等效加速，这就是通信税。
问：有没有现成的框架支持 3D 并行？
答：有，Megatron-LM 和 DeepSpeed 都支持，但配置复杂，需要深入理解并行策略。
问：3D 并行和模型量化、稀疏化有什么关系？
答：它们是正交的，你可以同时用 3D 并行和量化（如 FP8）来进一步降低显存，但会引入更多精度和速度的权衡。
能力边界：3D 并行不是银弹
你的模型能放进单卡，数据并行 + ZeRO 可能更高效；模型太大必须用 3D 并行，但你需要精心调优并行参数，并容忍训练速度的损失。如果你的硬件没有 NVLink 和高速 InfiniBand，3D 并行可能还不如用更小的模型。另外，3D 并行对训练软件的复杂性要求极高，调试一个 hang 住的通信死锁，能让你掉一层皮。
总的来说，3D 并行是大模型训练不得不走的钢丝，它让我们能训练千亿甚至万亿参数模型，但代价是极致的工程复杂度。如果你正准备上手，记住：先理解你的模型，再理解你的硬件拓扑，最后才去调并行策略——否则，你只会得到一堆 OOM 报错和慢如蜗牛的训练速度。
进阶阅读：激活重计算与3D并行
激活重计算（Activation Checkpointing）是3D并行的好搭档。它在前向时丢弃中间激活，反向时重新计算，用时间换空间。结合流水线并行，可以大幅减少显存占用，但会增加计算量。在 Megatron-LM 中，激活重计算是默认开启的。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/146.html