大规模训练：稳定性与分布式并行

训练需要走向分布式，通常只有两个原因：模型状态已经放不进一张加速器，或者只用一张加速器会耗时太久。分布式训练解决了这两个问题，也带来了新的问题。每个分片都必须和其他分片交换数据，每个参与训练的工作进程也都可能让任务停下来。

因此，设计分布式训练本质上是在做核算。面对任何一种布局，都要回答四个问题：

1. 哪些内容会复制，哪些内容会切分？
2. 每台设备的峰值显存是多少，其中是否计入临时缓冲区？
3. 一个训练步骤中，哪些消息位于关键路径上？
4. 为了在故障后正确续训，哪些状态必须保存下来？

本章从单设备上的一个训练步骤出发，逐步扩展到多维设备网格。这里没有普遍适用的配方。模型形状、序列长度、数值格式和网络拓扑会共同决定合适的布局。

从一个逻辑优化器步骤开始

数据并行是最直接的起点。$D$ 个副本各自接收批中的不同部分，运行同一个模型，再共同完成一次同步梯度更新。假设每个副本在优化器更新前处理 $m$ 个微批，每个微批包含 $B_\mu$ 条序列，那么

$$B_{\mathrm{global}} = DmB_\mu.$$

其中，$B_{\mathrm{global}}$ 是一次优化器步骤中的序列总数，$D$ 是数据并行度，$m$ 是每个副本累积的微批数，$B_\mu$ 是一个微批中的序列数。不同序列的词元数可能不同，因此生产训练通常同时记录每步序列数和每步词元数。

当各副本的本地批大小相同时，all-reduce 可以得到平均梯度：

$$\begin{aligned} g &= \frac{1}{D}\sum_{r=1}^{D}g_r, \\ g_r &= \frac{1}{B_{\mathrm{global}}/D} \sum_{i\in\mathcal B_r}\nabla_\theta \ell_i(\theta). \end{aligned}$$

这里，$r$ 是副本编号，$\mathcal B_r$ 是副本 $r$ 的本地批，$\theta$ 是参数向量，$\ell_i$ 是样本 $i$ 的损失，$g_r$ 和 $g$ 分别是本地平均梯度和全局平均梯度。局部批大小不同时，不能对各副本的均值做等权平均；归约时必须按样本数或词元数对各自的和加权。

这组方程规定了语义契约。不同切分方案可以改变张量存放的位置，也可以改变浮点归约的顺序，但它们都应实现同一个同步更新。位级一致是更强的要求，而且通常取决于拓扑。

核算模型状态显存

普通分布式数据并行（DDP）会在每个副本上复制参数、梯度和优化器状态。它切分了计算和批，却没有切分这些持久模型状态。零冗余优化器（ZeRO） 分三个阶段逐步消除复制 (Rajbhandari et al. 2020)：

其中，$P_\theta$ 是参数量；$b_w$、$b_g$ 和 $b_o$ 分别是工作参数、梯度以及全部优化器相关状态为每个参数占用的字节数；$D$ 是数据并行度。一种常见的 Adam 配方取 $b_w=2$、$b_g=2$、$b_o=12$：工作参数和梯度各占两个字节，FP32 主参数与两个 FP32 矩各占四个字节。激活和临时存储尚未计入时，这套配方需要 $16P_\theta$ 字节。数据类型是配方选择，并不是 Adam 的常数。

这张表给出的是持久状态估算，不是峰值显存估算。采用全分片时，计算某个封装单元之前要先聚合它的参数，计算之后再对梯度做归约与分散。一个有用的峰值下界是

$$\begin{aligned} M_{\mathrm{peak}} &\gtrsim \frac{P_\theta(b_w+b_g+b_o)}{D} + P_u b_w \\ &\quad + M_{\mathrm{act}} + M_{\mathrm{tmp}}. \end{aligned}$$

其中，$P_u$ 是当前单元与预取单元同时物化的参数量，$M_{\mathrm{act}}$ 是需要保存的激活显存，$M_{\mathrm{tmp}}$ 包括通信工作区、分配器开销、填充和其他临时缓冲区。PyTorch 全分片数据并行（FSDP） 则是同一全分片思路的原生版本：它把全分片应用到封装单元，并可在当前单元计算时预取后续单元的参数 (Zhao et al. 2023)。ZeRO-Infinity 进一步把内存层级扩展到 CPU 内存和 NVMe 存储 (Rajbhandari et al. 2021)。

切分得更多，并不会让通信消失。按 ZeRO 论文采用的元素传输量模型，DDP、ZeRO 第 1 阶段和第 2 阶段每步大约传输两份模型大小的数据；第 3 阶段约为三份，因为前向和反向都需要聚合参数 (Rajbhandari et al. 2020)。每设备持久状态随 $1/D$ 下降，但当 $D$ 增大时，每设备参数聚合载荷会趋近一份模型。更大的组仍可能因为延迟和拓扑而变慢。

另外四种切分工作的方式

数据并行只是一个轴。模型如果仍然放不下，或者单设备无法高效运行某些层，就需要继续切分。

张量并行：切开单层

张量并行（TP）会切分一层内部的矩阵乘法。对输入为 $X$ 的前馈块，一组配对的列切分与行切分可以写成

$$\begin{aligned} H_r &= \phi(XW_1^{(r)}), \\ Z &= \sum_{r=1}^{T}H_rW_2^{(r)}. \end{aligned}$$

其中，$T$ 是张量并行度，$r$ 是组内 rank，$W_1^{(r)}$ 是第一层投影的列分片，$W_2^{(r)}$ 是对应的第二层投影行分片，$\phi$ 是激活函数，构造 $Z$ 的求和需要一次归约。注意力也可以用同样思路切分查询头、键头和值头，再对输出投影做行切分。

Megatron-LM 的原始布局会安排好注意力块和前馈块，使每个 Transformer 层的前向传播需要两次 all-reduce，反向传播也有对应的归约 (Shoeybi et al. 2019)。TP 会减少每个 rank 的参数和算术量，却把通信放进每一层的关键路径。因此，它通常应映射到可用的最快互连域。这是一条放置原则，不是一条定律：TP 组能否跨节点，取决于消息大小、拓扑和性能目标。

流水线并行：切开层栈

流水线并行（PP）把连续的层组分配给不同阶段，再让微批依次通过这些阶段。阶段边界会向前发送激活，并在反向传播时送回激活梯度。阶段并不一定与节点边界重合。

同步刷新调度在开始和结束时会让一部分阶段空闲。假设各阶段耗时相同且通信可以忽略，GPipe 的理想空闲比例为 (Huang et al. 2019)

$$f_{\mathrm{bubble}}=\frac{p-1}{m+p-1}.$$

其中，$p$ 是流水线阶段数，$m$ 是每个副本在一次优化器步骤中处理的微批数，$f_{\mathrm{bubble}}$ 是填充与排空造成的空闲阶段时隙占比。分母不能省略：常见的 $(p-1)/m$ 是气泡时间相对于理想计算时间的比值，并不是气泡占总阶段容量的比例。

这个公式适合估算规模，不适合预测实测吞吐。阶段耗时不均、通信、优化器工作和其他调度都会改变结果。1F1B 调度可以减少保存的激活，但不会改变上述理想气泡；交错调度可以缩小气泡，却会发送更多点对点消息 (Narayanan et al. 2021)。PipeDream 研究的是异步调度和权重版本，它与 GPipe 的同步刷新属于不同的优化器语义取舍 (Narayanan et al. 2019)。

下面的可运行单元让气泡规模更直观。32 个阶段配 32 个微批时，即使在这个理想模型中，仍有接近一半的总阶段容量处于空闲状态。

def bubble_fraction(stages, microbatches):
    """均衡的 GPipe 刷新调度；忽略通信。"""
    if stages < 1 or microbatches < 1:
        raise ValueError("stages 和 microbatches 必须为正数")
    return (stages - 1) / (microbatches + stages - 1)


for stages in (4, 8, 16, 32):
    values = {
        microbatches: round(bubble_fraction(stages, microbatches), 3)
        for microbatches in (1, 4, 8, 16, 32, 64)
    }
    print(f"p={stages}: {values}")

print("p=32, m=32:", f"{bubble_fraction(32, 32):.1%}")

序列并行与上下文并行：切开词元

这两个名称描述的是相关但不同的机制。Megatron 风格的序列并行（SP）会沿序列维切分层归一化、dropout 以及其他原本由 TP 复制的区域。它复用 TP 组，把部分 all-reduce 换成 reduce-scatter 与 all-gather；因此，SP 通常不会成为设备网格里的另一个乘数 (Korthikanti et al. 2022)。选择性激活重算与 SP 相互补充，它只重算那些计算成本低、显存占用高的操作。

上下文并行（CP）会切分注意力计算中的序列。每个 rank 持有本地查询，再传递或聚合计算精确注意力所需的键值块。Ring Attention 使用在线 softmax，让键值块沿环传递 (Liu et al. 2023)；另一类实现会使用 all-to-all。CP 能减少每台设备保存的本地序列状态，但稠密注意力在全局仍要执行 $O(L^2)$ 算术，其中 $L$ 是上下文长度。最大上下文仍受总显存、通信、设备数和通信计算重叠程度限制。

专家并行：切开条件参数

混合专家模型还会引入专家并行（EP）。词元会被发送到持有所选专家的 rank，专家输出随后返回原来的词元所有者，常见实现采用 all-to-all (Lepikhin et al. 2020)。前向传播通常包含一次分发交换和一次合并交换，反向传播还要发送对应的梯度流量。热门专家会制造慢节点，因此路由局部性、负载均衡、可选的容量填充和无丢弃内核既是建模问题，也是系统问题。路由目标本身已在 第 9 章 讨论。

图 10.4 汇总了各条轴切分的状态和常见通信方式。图中的集合通信名称表示常见实现，并不是框架必须遵守的承诺；框架可以融合、拆解或重新调度这些操作。

把各条切分轴映射到网络

如果数据并行、张量并行、流水线并行和上下文并行是彼此独立的维度，那么稠密模型的设备网格满足

$$N = D T p C.$$

其中，$N$ 是总设备数，$D$ 是数据并行度，$T$ 是张量并行度，$p$ 是流水线阶段数，$C$ 是上下文并行度。SP 通常复用 $T$ 组。EP 取决于模型：它的组可能从其他网格维度中划出，也可能与其他组重叠，因此盲目再乘一个专家并行度会重复计算设备。

放置原则很直接：最频繁、对延迟最敏感的交换，应留在最快的互连域中。TP 常常优先占用这个域，因为每一层都会通信。EP 的 all-to-all 需要较强的对分带宽。PP 的消息大小取决于边界激活，ZeRO-3/FSDP 则要反复聚合参数。究竟让哪条轴跨越节点或机架，必须根据测量结果决定，而不能照搬固定层级 (Narayanan et al. 2021)。

不同集合通信的语义并不相同。all-reduce 会归约数值，并把结果交给每个 rank；reduce-scatter 会归约数值，但每个 rank 只保留一个分片；all-gather 会在每个 rank 上拼接所有分片；all-to-all 会向每个对端发送不同的分片。PP 主要使用点对点发送与接收。NCCL 等通信库会为这些操作提供多种算法，再根据硬件和消息形状进行选择 (NVIDIA 2024)。

通信是必要工作，但只有暴露在关键路径上的部分才会延长步骤时间。FSDP 可以在当前单元计算时聚合下一个单元；梯度桶可以在反向传播期间归约；PP 与 EP 消息也可以和相邻内核重叠。重叠并非免费，通信仍可能和计算争用显存带宽、执行资源或网络。要判断重叠是否生效，应看训练步骤轨迹，而不是只看配置中启用了哪些开关。

按操作选择数值格式

混合精度不是一种全局 dtype。常见配方会让大型矩阵乘法使用低精度格式，更谨慎地累积归约，并让敏感操作或优化器状态保留更宽的格式。最初的混合 FP16 配方使用 FP32 主参数、FP32 累积和损失缩放 (Micikevicius et al. 2017)。BF16 有八个指数位，正常数值范围大致与 FP32 相同，因此通常无需 FP16 那样的损失缩放；但它只有七个尾数位，累积和微小更新仍可能受精度限制。

FP8 常用两种布局：E4M3 把更多位用于精度，E5M2 把更多位用于范围 (Micikevicius et al. 2022)。数值必须先缩放到可表示区间，缩放粒度也是配方的一部分，可以按张量、分块或更小的块设置。Hopper 级硬件的 FP8 稠密矩阵乘法峰值可以达到 BF16 的两倍，FP8 操作数的载荷字节也只有一半；但类型转换、缩放元数据、非矩阵乘法操作以及保留的高精度状态，都会降低端到端收益。

DeepSeek-V3 展示了一次完整预训练：核心矩阵乘法使用 FP8，若干敏感操作和输出仍保留 BF16 或 FP32 (DeepSeek-AI 2024)。之后，NVIDIA 用 NVFP4 训练了一个 12B 模型，共处理十万亿词元，并报告了与 FP8 基线相近的结果。该配方使用细粒度缩放、Hadamard 变换、随机舍入和选择性高精度 (NVIDIA 2025)。这些都是经过实测的具体配方，不能证明所有模型或操作都能安全降到相同精度。

稳健的更新必须明确规定顺序：先累积完所有微批；如果启用了损失缩放，再解除 FP16 梯度缩放；跨所有 rank 检查数值是否有限；对逻辑上的完整分片梯度计算全局梯度范数；按配置裁剪；最后让所有分片一起更新或一起跳过更新。检查点还必须保存缩放状态，包括 FP8 的 amax 历史或动态损失缩放器。窄格式可能造成下溢、饱和、Inf/NaN、跳过步骤、发散或不易察觉的质量漂移，这些情况都必须验证。

吞吐测量必须写清假设

模型 FLOPs 利用率（MFU）把观测到的模型工作量与明确指定的硬件峰值作比较 (Chowdhery et al. 2023)：

$$\mathrm{MFU}= \frac{r_{\mathrm{tok}}F_{\mathrm{model,tok}}} {N P_{\mathrm{peak}}}.$$

其中，$r_{\mathrm{tok}}$ 是全局实测每秒词元数，$F_{\mathrm{model,tok}}$ 是模型对每个词元执行前向与反向传播所需的 FLOPs，$N$ 是加速器数量，$P_{\mathrm{peak}}$ 是单台加速器在对应精度下的峰值 FLOP 率。报告必须说明采用的 FLOP 口径、如何处理 MoE 的活跃参数和填充、使用哪一种精度峰值，以及测量窗口多长。

硬件 FLOPs 利用率（HFU）与 MFU 的口径不同：HFU 还会计入激活重算等实际执行的 FLOPs，因此即使有效词元吞吐没有变化，HFU 也可能高于 MFU。两个指标都不能证明训练质量或可靠性。稳态 MFU 通常不计重启时间，端到端有效吞吐还应计入失败步骤、检查点、评测暂停和输入停顿。

MFU 偏低时，应先检查轨迹，再调整网格。常见原因包括内核过小或形状不佳、阶段不均衡、暴露通信、重算过多、输入停顿，以及同步等待慢节点。峰值显存、各通信组的集合通信耗时、每秒词元数、步骤时间尾部和恢复损失时间，才构成完整的运行视图。

把正确重启纳入正确性要求

普通同步任务中，只要一个 rank 失效，它所在的通信组就不再有效。弹性启动器通常会停止并重建整个工作进程组，再加载已经提交的检查点；它们不会把一台设备热替换进仍在运行的集合通信。完整控制流程是：检测训练不再前进；中止通信；隔离可疑工作进程；补充容量；重新会合；重建通信组；加载经过验证的检查点；核对数据游标；然后恢复训练。

可恢复的检查点是一个原子集合，不只是一份参数文件。它要包含模型参数和任何主参数、优化器张量与步骤、学习率状态、损失缩放器和低精度缩放历史、各 rank 的随机数状态、采样器与打乱状态、数据混合计数与游标，以及分片清单、校验和与完成标记。在优化器步骤边界保存，可以避免记录只累积了一部分的梯度。还应保留不止最近一次检查点，因为数值已经损坏的状态也可能被保存下来。

恢复有两种合理契约。确定性重放力图再次使用相同批次和随机数，通常要求拓扑不变，并使用确定性内核。覆盖等价恢复可以在不同总设备数上重新切分，只要整体消费的数据保持一致；它不要求位级相同。必须明确选择哪一种契约，并对它进行测试。

检查点频率是在写入成本和重复计算之间取舍。一个一阶的故障停止模型给出

$$\begin{aligned} W(\tau)&\approx\frac{C}{\tau}+\frac{\tau}{2M}, \\ \tau^\star&\approx\sqrt{2CM}. \end{aligned}$$

其中，$W(\tau)$ 是检查点与重算造成的预期时间损失比例，$\tau$ 是两次检查点之间的有效计算时间，$C$ 是暴露在关键路径上的检查点成本，$M$ 是同步任务的平均中断间隔，$\tau^\star$ 是近似最优间隔 (Daly 2006)。这一近似假设故障停止事件相互独立，并省略了相关故障、恢复延迟和存储争用。异步检查点会降低暴露的 $C$，但仍会消耗设备链路、主机内存、网络和存储带宽。CheckFreq 展示的是自适应频率与流水化快照、持久化阶段，并不是免费的写入 (Mohan et al. 2021)。

长跑前验证布局

只有通过小规模而有针对性的测试，分布式配置才适合进入长跑：

• 更新等价性： 在一个很小的确定性批上，用明确容差比较未切分与已切分实现的损失、梯度和一次优化器更新。
• 显存核算： 在不同微批大小和序列长度下测量持久、峰值和临时显存，确认封装单元与预取假设和轨迹一致。
• 通信核算： 分析每个通信组，区分暴露通信与总通信，并测试计划中的跨节点边界。
• 数值稳定性： 用更高精度对照比较逐操作精度配方，记录非有限张量、跳过步骤、裁剪前梯度范数、发生裁剪的步骤比例和评测漂移。
• 重启正确性： 在计算期间和检查点持久化期间分别终止一个工作进程，拒绝不完整清单，恢复声明的数据契约，并把恢复后的轨迹与对照比较。
• 有效吞吐： 运行时间要足以覆盖检查点、输入加载和至少一次恢复，并在 MFU 之外报告每墙钟秒的有效词元数。

FlashAttention 展示了为什么每一层核算都不可省略。它不在 HBM 中物化完整分数矩阵，也能计算精确注意力 (Dao et al. 2022)。设批大小为 $B$、头数为 $H$、上下文长度为 $L$、每个分数元素占 $b$ 字节，仅物化分数就需要

$$M_{\mathrm{score}} = BHL^2b.$$

其中，$M_{\mathrm{score}}$ 不包括 softmax 工作区、输入、输出、保存的激活和分配器开销。FlashAttention 去掉了这个随 $L^2$ 增长的 HBM 张量，但不会消除稠密注意力的算术量，也不会取消对超长上下文做切分的需要。

整次训练都应遵守同一种核算方法：说清张量名称，计算它占用的字节，指出它何时存在，再说明后续需要哪一次通信或恢复操作。这样，各种并行缩写才会变成一份能够验证的工程设计。

延伸阅读