08. 分布式训练：把数据、状态、层和序列分到多张卡

第 3 篇算过，Llama-3.1-8B 全参训练时，仅权重、梯度和优化器状态就要约 128.5 GB，一张 80 GB 的 H100 放不下；第 4 篇算过，把它从头预训练一遍至少要几十万 GPU 小时，一张卡要算几十年。多卡训练要解决的就是这两件事：放不下和算不完。这一篇从资料助手团队手里的一台 8 卡服务器讲起，逐个说明数据并行、ZeRO/FSDP、张量并行、流水线并行、序列与上下文并行、专家并行各自切开了什么、要付出多少通信，再看 Llama 3 405B 怎样把它们组合到一万六千张卡上，以及规模大到这个程度时，故障怎样成为日常。读完你应该能为一个给定的模型和机器数，说出该先用哪种切法、瓶颈会出在哪里。

推理也会用多卡（推理系列第 16 篇），但训练要多管两样东西：反向传播产生的梯度，以及比权重大好几倍的优化器状态。所以训练的切法更多，通信也更重。

先看全貌：六种切法

同样 4 张卡，不同的并行方式切的是不同的对象。下图每张卡里的四条横条代表模型的 4 层，涂色部分是这张卡常驻保存的参数（以及对应的梯度和优化器状态），下方灰条是这张卡处理的数据：

数据并行和 ZeRO/FSDP 切的是“谁处理哪些样本、谁保存哪部分状态”；张量并行和流水线并行切的是模型本身；上下文并行切的是一条样本内部的序列；专家并行只用于 MoE 模型。下面按资料助手团队遇到问题的顺序逐个展开。

数据并行：每张卡一份完整副本

数据并行（Data Parallelism，DP）最直观：8 张卡各放一份完整的模型，每张卡拿不同的样本做前向和反向，各自算出一份梯度。更新参数之前，8 份梯度要求平均，让每张卡拿到同一个结果，然后各自执行同样的更新，参数始终保持一致。这个“各自求和、人人拿到结果”的集合通信操作叫 all-reduce。效果相当于把批大小扩大了 8 倍。

通信量可以手算。常用的环形 all-reduce 中，每张卡要收发约 2 × (n−1)/n 倍的数据量：

8B 模型的 BF16 梯度：8.03×10⁹ × 2 字节 ≈ 16.1 GB
8 卡环形 all-reduce，每卡收发 ≈ 2 × 7/8 × 16.1 ≈ 28 GB / 步

机内 NVLink（H100 合计 900 GB/s，单向约 450 GB/s）：约 0.06 秒
机间 400 Gb/s 网卡（约 50 GB/s）：约 0.56 秒

这笔通信每一步只发生一次，和这一步处理多少 token 无关。如果每张卡一步处理 12.5 万个 token，计算量约 6 × 8.03×10⁹ × 1.25×10⁵ ≈ 6×10¹⁵ 次，按每卡 400 TFLOPS 约 15 秒，通信只占零头；而且框架会把梯度分桶，反向算完一桶就开始通信一桶，大部分通信可以和计算重叠。所以数据并行的扩展效率通常很好。

它的问题是不省显存：每张卡仍然要放下完整的 128.5 GB 模型状态。资料助手做 LoRA 时，冻结的底座约 16.1 GB、可训练部分约 0.67 GB，单卡放得下，直接用 8 卡数据并行把 3 轮训练从一两个小时缩到十几分钟就够了。想做 8B 全参训练，就需要下一种方法。

ZeRO 与 FSDP：把重复保存的状态切开

数据并行里，8 张卡保存着 8 份一模一样的权重、梯度和优化器状态。ZeRO 论文（Rajbhandari 等，2019）的出发点就是消除这种冗余：既然每张卡最终都要做同样的更新，那就让每张卡只负责 1/8 的参数，只保存这 1/8 对应的状态。按切分的深度分三个阶段。用 Ψ 表示参数量，n 表示卡数：

每参数 16 字节 = BF16 参数 2 + BF16 梯度 2 + FP32 主权重、一阶动量、二阶动量 12

普通数据并行：16Ψ                   = 128.5 GB
ZeRO-1 切优化器状态：4Ψ + 12Ψ/n     = 32.1 + 12.0 = 44.2 GB
ZeRO-2 再切梯度：  2Ψ + 14Ψ/n       = 16.1 + 14.1 = 30.1 GB
ZeRO-3 再切参数：  16Ψ/n            = 16.1 GB

（Ψ = 80.3 亿，n = 8；只算模型状态，激活另计）

ZeRO-1 和 ZeRO-2 几乎不增加通信。普通数据并行的 all-reduce 本来就可以拆成两半：先 reduce-scatter（求和后每张卡只拿到自己负责的那 1/n），再 all-gather（把各卡更新好的那份参数拼回完整的一份），总量仍是 2Ψ。ZeRO 只是把“拿到 1/n 梯度后更新 1/n 参数”放到两半中间做。ZeRO-3 连参数也只存 1/n，每一层计算前要临时把这一层的完整参数收集过来，算完即丢；反向时还要再收集一次，于是通信变成 3Ψ，是普通数据并行的 1.5 倍。论文里的例子是 7.5B 模型、64 卡，每卡模型状态从 120 GB 依次降到 31.4 GB、16.6 GB、1.88 GB。

PyTorch 的 FSDP（Fully Sharded Data Parallel，完全分片数据并行）是同一思路的原生实现，完全分片模式对应 ZeRO-3，只分片梯度和优化器状态的模式对应 ZeRO-2；DeepSpeed 则直接以 ZeRO 的阶段命名。对资料助手团队来说，8B 全参 SFT 在一台 8 卡 H100 上用 ZeRO-3 或 FSDP，每卡模型状态约 16.1 GB，再加上开了激活重计算的激活，就能放下。

ZeRO 切的是模型状态，不切激活：每张卡仍然要为自己那批样本保存完整的激活。另外，ZeRO-3 每步要把全部参数收集两遍，网络慢时会成为瓶颈。跨多台机器时常用“混合分片”：机内 8 卡之间分片，机器之间按普通数据并行复制，把大流量留在机内高速链路上。

张量并行：把每一层的矩阵切开

当模型大到连一层的计算都希望分给多张卡，或者激活太大放不下时，就要切模型本身。张量并行（Tensor Parallelism，TP）来自 Megatron-LM（Shoeybi 等，2019），它把每一层里的矩阵切开，每张卡算一部分。

以 Llama-3.1-8B 的前馈层为例，它先把 4096 维的向量升到 14336 维，经过激活函数，再降回 4096 维。张量并行在 8 张卡上这样切：

• 第一个矩阵按列切：每张卡负责 14336/8 = 1792 个输出维度。每张卡都拿到完整的输入，各自算出中间结果的 1/8。激活函数逐元素计算，各卡对自己那份直接算，不需要通信。
• 第二个矩阵按行切：每张卡用自己那 1792 维的中间结果，乘以对应的 1792 行，得到一个完整 4096 维、但只包含部分和的输出。
• 最后做一次 all-reduce，把 8 份部分和加起来，得到真正的输出。

注意力层也是同样的套路：32 个查询头、8 个 KV 头按头分给 8 张卡，每卡 4 个查询头、1 个 KV 头，最后的输出投影按行切，再 all-reduce 一次。这样每一层前向要 2 次 all-reduce，反向还要 2 次。

这些通信夹在每一层的计算中间，后面的计算必须等它完成，很难重叠。粗算一下：

每次 all-reduce 的数据：一个微批次 8192 token × 4096 维 × 2 字节 ≈ 67 MB
每卡收发：2 × 7/8 × 67 MB ≈ 117 MB；每层 4 次 × 32 层 = 128 次 ≈ 15 GB
每卡计算：6 × 8.03×10⁹ × 8192 ÷ 8 卡 ≈ 4.9×10¹³ 次，按 400 TFLOPS 约 0.12 秒

走 NVLink（约 450 GB/s）：约 0.03 秒，是计算时间的四分之一
走机间网络（约 50 GB/s）：约 0.3 秒，比计算还长

以上只是带宽估算，没算延迟。结论已经清楚：张量并行必须放在 NVLink 连接的同一台机器内，并行度一般不超过一台机器的卡数 8。

张量并行里还有一些操作（LayerNorm、Dropout）是每张卡重复计算、重复保存激活的。Megatron 团队 2022 年提出的序列并行把这些操作沿序列方向切给各卡，配合选择性重计算，论文报告激活显存减少约 5 倍。

流水线并行：按层分段接力

流水线并行（Pipeline Parallelism，PP）按层切：32 层分给 4 张卡，每张 8 层。第一张卡算完前 8 层，把激活交给第二张卡，依此类推；反向时梯度沿相反方向传回。卡与卡之间只传一个微批次的激活和激活的梯度，通信量小，而且是点对点传递，可以跨机器。

代价是等待。如果一次只送一批数据，第二张卡要等第一张卡算完才有活干，大部分时间都在空等。解决方法是把一批数据切成多个微批次，让它们像流水线上的零件一样依次进入。但开头填满流水线、结尾排空流水线的时候，总有卡在空等，这部分叫气泡：

图里前向占 1 格、反向约占 2 格。4 段流水线、8 个微批次时，总时长 33 格，每张卡空等 9 格：

气泡占比 = (p − 1) / (m + p − 1) = 3 / 11 ≈ 27%
p：流水线段数；m：每批的微批次数

增加微批次数可以压低气泡：m = 32 时降到 3/35，约 9%。上面一种调度（GPipe）先做完所有前向再做反向，第一张卡要同时保存 8 个微批次的激活。下面的 1F1B 调度（一个前向、一个反向交替）先做几个前向把流水线填满，之后每完成一个前向就做一个反向，及时释放激活。气泡大小不变，但第一张卡最多只保存 p = 4 个微批次的激活，微批次再多也不会撑爆显存。

再进一步的交错式调度让每张卡负责几段不相邻的层（例如第 1、5 段），气泡能再缩小到原来的几分之一，代价是更多的点对点通信。Megatron 团队 2021 年的论文用 TP、PP、DP 组合配合交错调度，在 3072 张 A100 上训练 1 万亿参数的模型，每卡达到峰值的 52%。

流水线还要注意各段的负载均衡：第一段多了词嵌入，最后一段多了输出层和损失计算，比中间段更重。Llama 3 的做法是让首尾两段各少放一层。

长上下文：上下文并行

训练 128K 长上下文时，一条样本的激活本身就可能超出单卡显存。上下文并行（Context Parallelism，CP）把同一条序列切成几段，每张卡处理一段。前馈层对每个 token 独立计算，各段互不干扰；注意力层每个 token 需要看到前面所有 token 的 K、V，于是各卡要交换 K、V，常见做法是沿环形依次传递，或者直接 all-gather 收集全部 K、V。Llama 3 在长上下文阶段用了 16 路上下文并行，序列长度 131,072。

它和上一节 Megatron 的序列并行名字相近，切的范围不同：序列并行只切张量并行组里那些逐 token 的操作，上下文并行连注意力也一起切。

MoE：专家并行

MoE（混合专家）模型每层有很多个前馈“专家”，每个 token 只交给其中几个计算。专家并行（Expert Parallelism，EP）把不同的专家放到不同的卡上，token 按路由结果通过 all-to-all 通信发到专家所在的卡，算完再发回来，每个 MoE 层两次。它的难点是负载不均：如果路由总把 token 送给少数几个热门专家，这几张卡忙不过来，其他卡空等。训练时通常要加辅助损失或其他机制让路由更均衡。推理侧的拆法见推理系列第 16 篇。

组合起来：Llama 3 405B 怎样用一万六千张卡

单一的切法都有上限：张量并行出不了一台机器，流水线段数越多气泡越大，数据并行不省显存。大模型训练把它们按通信密度嵌套起来：通信最密、最怕延迟的放在最内层，走最快的链路；通信最稀、能和计算重叠的放在最外层。

Llama 3 论文（2024）按 [TP, CP, PP, DP] 的顺序组织 405B 的预训练：

阶段	卡数	TP	CP	PP	DP	序列长度	每卡算力	MFU
预训练	8,192	8	1	16	64	8,192	430 TFLOPS	43%
预训练	16,384	8	1	16	128	8,192	400 TFLOPS	41%
长上下文	16,384	8	16	16	8	131,072	380 TFLOPS	38%

每批都是 1600 万 token。张量并行 8 路，正好是一台服务器的 8 张卡；16 段流水线跨 16 台服务器，构成一个完整的模型副本，共 128 张卡；数据并行用 FSDP 实现，分片模型、梯度和优化器状态。粗算一下显存：405B 按每参数 16 字节约 6.48 TB，被 TP 和 PP 切成 128 份后每卡约 50.6 GB，FSDP 再在数据并行组里分片，剩下的空间才能留给激活。

表里还能看出两个规律。卡数从 8192 翻倍到 16384，每卡算力从 430 降到 400 TFLOPS，规模越大，通信和同步的损耗越多。进入长上下文阶段后，序列长了 16 倍，数据并行的份数从 128 降到 8，把卡让给上下文并行，MFU 又降到 38%。

资料助手团队用不到这么复杂的布局，但选法是一样的。一台机器能放下时，用 ZeRO/FSDP 加数据并行最简单；模型大到一台机器放不下（例如 70B 全参训练），先在机内加张量并行，或者跨机用 FSDP 分片；再大，才需要流水线。

规模越大，越要为故障设计

一张卡一年出一次故障，听起来很可靠；一万六千张卡放在一起，故障就是每天都会发生的事。Llama 3 论文统计了预训练期间一段 54 天的记录：

• 共 466 次作业中断，其中 47 次是计划内的（升级固件、调整配置等），419 次是意外中断；
• 意外中断中约 78% 被确认或怀疑是硬件问题，次数最多的两类是 GPU 故障（148 次）和 GPU 的 HBM3 显存故障（72 次）；
• 尽管如此，有效训练时间仍超过 90%，期间只有 3 次需要较大的人工介入；
• 甚至能观察到 1%—2% 的昼夜吞吐波动，与午间气温影响 GPU 动态调频有关。

419 次意外除以 54 天，平均每天将近 8 次，约每 3 小时一次。同步训练里所有卡要一起完成每一步，任何一张卡掉线，整个作业都会停下，只能从最近的检查点重启。

所以大规模训练的核心工程问题之一是检查点：定期把模型参数、优化器状态、数据读取位置和随机数状态完整保存下来。它的规模也不小，405B 按每参数 16 字节算，一次完整检查点约 6.48 TB，要让每张卡并行写自己负责的分片，最好异步进行，不阻塞训练。保存间隔是一个权衡：太频繁，大量时间花在写存储上；太稀疏，每次故障丢掉的进度更多。如果平均 3 小时出一次故障、每 30 分钟存一次，平均每次故障损失约 15 分钟的计算，再加上重启和加载的时间。

小规模训练同样需要检查点，只是理由不同：云上的抢占式实例可能随时被回收，几个小时的训练跑到一半断掉很常见。断点续训最容易出的问题是只保存了模型参数，没保存优化器状态和数据位置，恢复后 loss 曲线出现一个跳变，或者某些数据被重复训练。恢复后先对比续训前后几十步的 loss，确认曲线连续（第 22 篇讲断点续训的工程细节）。

多卡训练按切的对象分成几类：数据并行切样本，ZeRO/FSDP 切重复的模型状态，张量并行切每层的矩阵，流水线并行切层，上下文并行切序列，专家并行切 MoE 的专家。每种切法都用通信换显存或时间，通信越密的越要放在越快的链路上，于是有了“TP 在机内、PP 跨机、DP 在最外层”的布局。规模越大，故障越频繁，检查点和自动恢复决定了实际能用上多少算力。资料助手用一台 8 卡机器就能完成 8B 的全参或 LoRA 微调，接下来回到训练内容本身，看怎样准备数据、让模型学会按公司口径回答：指令微调 SFT。