第2章分布式训练基础

分布式训练基础：数据并行、模型并行、张量并行的基本概念

聊分布式训练之前，我先问个问题：
你手头有一张A100，显存80G，但模型参数有200G。怎么办？
答案很简单——用多张卡。但怎么用，这里头门道就多了。

我个人习惯把分布式训练分成三大流派：数据并行、模型并行、张量并行。
说白了，就是三种切分思路。你想想看，一个模型无非就是数据和参数，怎么分？

数据并行，我估计你最早接触的就是它。
每张卡上都放一份完整的模型副本，然后把训练数据切成多份，每张卡处理一份。

举个例子：
你有4张卡，一个batch有64条数据。每张卡分到16条，各自算梯度。
算完之后，大家把梯度汇总一下，更新模型参数。

嗯，这里要注意：
数据并行最大的坑是通信开销。每张卡算完梯度后，得把所有梯度同步一遍。
我在项目中遇到过，模型参数特别大时，梯度同步的时间比计算时间还长。

  核心公式：

  总吞吐 ≈ 单卡吞吐 × 卡数 / (1 + 通信开销占比)

说白了，卡越多，通信越频繁。如果通信带宽跟不上，加卡反而没收益。

模型并行，就是模型太大，一张卡放不下。
把模型的不同层放到不同卡上。比如，前几层在卡0，中间层在卡1，后几层在卡2。

这样做的好处是：每张卡只存一部分参数，显存压力小。
但坏处也很明显：计算是串行的。卡0算完才能传给卡1，卡1算完才能传给卡2。
你想想看，这就像流水线，只要有一张卡慢，整条线都等着。

我的经验：
模型并行适合那种层数特别深的模型，比如Transformer的几十层。
但要注意，层与层之间的通信量很大，尤其是激活值传递。

我曾经踩过一个坑：
把模型并行和数据并行混着用，结果通信拓扑搞得太复杂，调试了三天才跑通。
后来我建议，能不用模型并行就别用，除非模型真的放不下。

张量并行，比模型并行更细。
它不是按层切，而是把一个层内部的矩阵运算拆开。

比如一个全连接层，权重矩阵是 [4096, 4096]。
我可以把它切成两块 [4096, 2048] 和 [4096, 2048]，分别放在两张卡上。
输入数据也切成两半，各自算一部分，最后拼起来。

这样做的好处是：
每张卡只算一半的矩阵乘法，计算量减半，显存也减半。
但坏处是：通信更频繁。每次矩阵乘法前后都要做 all-reduce 或 all-gather。

注意：
张量并行对通信带宽要求极高。我建议只在单机多卡场景下使用，跨机器的延迟太高，得不偿失。

实际训练大模型时，这三种方式经常混着用。
比如训练GPT-3，用了数据并行 + 张量并行 + 模型并行三层嵌套。

我画了一张图，帮你理解它们的关系：

从这张图你能看出来：
数据并行在最外层，负责数据分发和梯度汇总。
模型并行在中间层，负责把模型按层切分。
张量并行在最内层，负责把每一层的计算拆得更细。

我的建议：
刚开始做分布式训练，先从数据并行入手。
等遇到显存瓶颈了，再考虑加张量并行。
模型并行，我建议最后再碰，调试成本太高。

嗯，这三种并行方式，说白了就是空间换时间和时间换空间的权衡。
数据并行用多份模型换训练速度，模型并行和张量并行用通信换显存。
怎么选，取决于你的硬件条件和模型大小。

我个人习惯，先算一笔账：
模型参数多大？单卡显存多少？通信带宽多少？
算清楚了，再决定用哪种并行方式，或者怎么组合。

最后说一句：
分布式训练不是银弹。有时候加卡反而更慢，就是因为通信开销超过了计算收益。
这个平衡点，就是我们下一节要聊的——通信与计算的权衡。

公众号:蓝海资料掘金营，微信deep3321