第 27 章 大规模训练：系统分工与资源估算

27.1 模型装得下，一步训练却未必装得下

假设一百万个权重占两百万字节，设备还有一千万字节可用。看起来训练应该绰绰有余。但第 19 章的更新需要梯度，优化器也会保留跨步骤使用的信息；第 20 章还要保存中间值，让反向计算能沿着前向运算走回去。这些对象都没算在两百万字节的权重文件里。

本章先为一个给定的小模型列资源账，再讨论怎样分工。整个练习只需 CPU 做标量运算，不要求 GPU，也不表示本书已经运行工业训练集群。MG-2 已有真实的训练与恢复实例，但它使用 float64，其内存账不能套用下面的混合精度设定。

27.2 持续保留的数组，以及运算之间的中间值

例 27.1：一百万参数的存储账。 给定 $N=1{,}000{,}000$ 个可训练参数。工作权重和累积梯度，每个参数分别占 2 字节；另保留一份每个参数 4 字节的高精度主副本。Adam 使用两个各占 4 字节的矩数组：一个跟踪梯度的移动平均，另一个跟踪梯度平方的移动平均。这些是明确给定的存储假设，并非所有优化器、框架都如此。

表 27.1：按设定推导的模型状态存储量

对象	为什么需要保留	每个参数对应字节	总字节
工作权重	用于前向与反向计算	2	2,000,000
梯度	更新前汇总各项贡献	2	2,000,000
主权重副本	累积更细小的更新	4	4,000,000
一阶矩与二阶矩	保留 Adam 的更新历史	4 + 4	8,000,000
合计	本估算中以上对象同时存在	16	16,000,000

主副本与两个矩数组合计，每个参数占 12 字节。有些实现采用不同精度或副本安排；多算一个不存在的副本，与漏算一个实际存在的副本，同样不对。按这里的设定，激活还没加进去，一千万字节的设备就已经装不下了。

再给定本地微批次包含两个序列，每个序列 128 个 token 位置，宽度 64，共四层。为便于起步，假设每层保留十个形状均为“批次 × 位置 × 宽度”的数组，每个值占 2 字节。激活存储量为：

$$M_{\mathrm{act}}=2\times128\times64\times4\times10\times2=1{,}310{,}720\text{ bytes}.$$

各因子依次表示序列数、位置数、坐标数、层数、每层保留的数组数，以及每个坐标的字节数。激活（activation） 是前向计算产生的值，不是另一套待训练参数。一层中的一个数组占 $2\times128\times64\times2=32{,}768$ 字节，四十个这样的数组合计 1.25 MiB。

模型状态与已保存激活合计 17,310,720 字节。若再预留给定的 25% 余量，就得到 $17{,}310{,}720\times1.25=21{,}638{,}400$ 字节。这是规划余量，不是测得的内存分配峰值。有些实现会保存随序列长度平方增长的注意力分数；嵌入、logits、临时工作区、通信缓冲区和内存碎片也要另算。“十个数组”的假设不会自动包含这些对象。要承诺某设备能够承载训练，最终还须用真实分配记录取代近似。

为模型状态保留的存储

1. 模型保存的数值工作权重 2 MB + 主副本 4 MB
2. 一组数字梯度 2 MB + 矩状态 8 MB
3. 处理步骤主副本／矩状态分给 4 个进程

仅分片优化器相关状态，每个进程的这部分存储从 16 MB 降至 7 MB；梯度在更新边界重置。

单步计算内

1. 一组数字40 个已保存数组 × 32,768 字节
2. 处理步骤保留 16 个数组，其余重算
3. 一组数字1,310,720 → 524,288 字节

重算减少激活存储，增加前向计算。

跨停止恢复

1. 模型之外一致的权重与优化器状态
2. 模型之外步数 + 随机状态 + 数据游标 + 版本

在同一更新边界恢复所需状态；仅有权重文件不能复现原来的继续过程。

图 27.1：不同存储对象有不同存活时间。 本书原创图示 · CC BY-SA 4.0。

图 27.1 区分了跨更新保留的状态、可以丢弃或重算的激活，以及恢复运行所需的信息。它们都占存储，却有不同的存活时间。下面的节省办法正是从这些差异中产生的。

27.3 缩小数值表示，还要保住更新

工作权重已经用了较小的表示，为什么还要保留主副本？先用十进制舍入作类比。一个只存两位小数的工作权重为 1.00，每次减去 0.003 后立即舍入，结果仍为 1.00。若始终只保留舍入后的副本，每次更新都会丢失。若较精细的主副本先保留 0.997，再保留 0.994，第二次把它舍入给计算使用，就得到 0.99。微小更新先积累起来，才反映到工作副本中。

实际的 FP16 和 BF16 使用二进制浮点格式，不按这个十进制规则存储。两者虽然都占 2 字节，却在有效精度与指数范围之间作了不同分配。因此，“字节减半”不表示溢出、下溢和舍入行为相同。混合精度（mixed precision） 要指定哪些运算、哪些存储数组使用哪种表示，并非给所有张量统一换个名字。

梯度也可能小到无法表示。损失缩放（loss scaling） 先放大损失，让反向计算产生相应放大的梯度，再在更新前去掉这个倍数。例如缩放倍数为 1,024 时，0.000001 的梯度会在缩放后的计算中变为 0.001024；随后在足够精度下除以 1,024，恢复原本的幅度。已经在缩放之前丢掉的数值无法因此找回，倍数过大又可能导致溢出。训练循环须检查非有限值，并一致地处理跳过更新时的调度。MG-2 的 float64 实验展示了更新顺序，没有实现 FP16 损失缩放。

27.4 几个小批次，可以共同决定一次更新

在本章线性估算下，减小微批次可以减少激活存储，但也会改变一次更新平均了哪些证据。梯度累积（gradient accumulation） 把这两件事拆开：保持参数不变，依次处理若干微批次，累积贡献，最后只更新一次。

例 27.2：分母应当数有效预测事件。 两个微批次分别包含 2 个和 6 个有效目标 token 事件。对某个参数，它们的平均梯度分别为 1 和 3，因此梯度总和分别是 $2\times1=2$ 和 $6\times3=18$。八个事件合起来，平均梯度为 $(2+18)/8=2.5$。若直接平均两个已报告的均值，就得到 $(1+3)/2=2$，把较小微批次的影响放大了。

第 i 个微批次有 $n_i$ 个有效事件，平均梯度为 $g_i$，则合并后的均值为：

$$\bar g=\frac{\sum_i n_i g_i}{\sum_i n_i}.$$

分母数的是真正贡献损失的事件，不包括被 Mask 排除的填充位置；分子把各微批次的梯度总和还原出来。各微批次的有效事件数相同时，把每份平均损失除以微批次数就够了；有效事件数不同时，必须按这个数量加权。梯度裁剪和优化器更新应放在约定的合并更新边界。如果第一个微批次后就改参数，第二个梯度的计算起点也变了，这就不再是同一种过程。

例如，一次更新要汇总八个等长序列，可以在一个工作进程中使用本地批量 2、累积 4 次；也可以让两个数据并行工作进程各用本地批量 2、累积 2 次。长度不等时，还应公开实际汇总的有效 token 数。累积减少了同时保存的激活，但计算仍要做，并没有凭空增加计算设备。

27.5 把数据、矩阵和层分给不同进程

这里的 工作进程（worker） 指参与训练协作的一个进程。数据并行（data parallelism） 让不同进程处理不同数据，同时从一致的模型参数出发。各自计算梯度贡献后，再通过集体通信汇总，使更新保持一致。如果损失按 token 求均值，汇总也必须遵守例 27.2 的分母规则。

表 27.2：分工方式不同，跨边界传递的东西也不同

分工	分开哪部分	边界上需要传什么
数据并行	训练样例	梯度总和／计数，以及一致的更新状态
张量并行	同一个矩阵或层的不同部分	完成该层运算所需的局部结果
流水线并行	连续的不同层	向前传激活，向后传相应梯度
状态分片	优化器、梯度或权重状态的副本	已汇总的梯度分片，或当前运算所需的权重

张量并行（tensor parallelism） 可以让一个进程算矩阵的部分输出列，另一个算剩余列；后续运算决定这些结果需要拼接还是相加。流水线并行（pipeline parallelism） 把前后不同层放到不同进程。微批次能让各阶段交错工作，但流水线填充和排空时，某些阶段可能闲置。调度还必须规定何时更新权重。图上画了四个方框，不足以证明速度变成四倍。

例 27.3：先分掉最大的副本。 用四个工作进程训练例 27.1 的模型。先只把一千二百万字节的主副本／矩状态平均分片，每个进程保留三百万字节；再加完整的两百万字节工作权重和两百万字节梯度，每个进程共七百万字节。计入给定的本地激活后为 8,310,720 字节，尚未计缓冲区和余量。

若梯度也按理想情况分片，每个进程的模型状态变为 $2+14/4=5.5$ MB。若全部 16 MB 模型状态都理想分片，长期保留的份额为 4 MB。这是逐步改变的不同安排，不能随意把某一行数字代到另一种方案里。执行某层时可能需要临时收集权重，峰值会超过长期保留的份额。进程本地不再保存的信息，现在要靠通信提供。

假设一次传输需搬运给定的一百万字节，链路每秒实际传一亿字节。仅负载本身就需 $1{,}000{,}000/100{,}000{,}000=0.01$ 秒，即 10 ms，尚未计延迟和争用。如果有用计算只需 2 ms，传输又不能与它重叠，这次迭代至少需要 12 ms。增加进程可能把瓶颈从算术移到通信。

27.6 省下激活后怎样重算，故障后怎样恢复

激活检查点（activation checkpointing） 只保留选定的中间激活，在反向计算时重新算出缺失部分。若把本章每层保留的数组从十个减至四个，保存的激活就降为 524,288 字节。代价是多做一些前向运算来重建省去的值。它用计算换空间，不会减少表 27.1 的权重和矩状态。

恢复检查点（recovery checkpoint） 则服务于另一个目的：停止后接着训练。第 21 章 MG-2 已保存模型与优化器状态、完成步数、采样／随机状态和配置。数据游标或批次生成器状态不能省：恢复权重后如果换了后续样例，就不是原来那条训练过程的继续。分片训练还需记录分片身份、一致的更新边界、tokenizer／数据版本，并确认所有分片属于同一次保存。只写了一部分的文件不能当作完整检查点。

小规模验收可以比较两条路径：一条连续更新，另一条保存、重载后继续，两者使用相同且明确的 CPU 环境。MG-2 已实现这个检查；本部分不声称新增了分布式恢复实验。跨硬件或软件变化后仍逐位相同，需要另行举证。可加载的权重文件足以支持某些模型运行用途，却不够满足这里的训练恢复契约。

现在，资源估算表可以说明自己算了什么、漏了什么、分片或重算的是哪类数组，以及新增哪些传输。它仍不能证明工业训练能力。进一步改模型之前，我们也能提出另一类问题：怎样的证据才能表明某个内部组件负责某种行为？

27.7 本章小结

训练内存包括工作权重、梯度、优化器／主副本状态和激活，各自存活时间不同。精度、累积、分片、并行和重算针对不同项；它们必须保留预期的更新含义，或者明确说明改变了什么。恢复训练需要保存过程状态，不能只留下最新权重。

27.8 习题

1. 保持例 27.1 的参数状态不变，把微批次加倍为四个序列。按十个数组的假设，计算激活字节数及加余量前的总量。
2. 两个微批次分别有 3 个、1 个有效事件，平均梯度为 2、−2。求合并梯度，并解释等权平均两个均值错在哪里。
3. 有人把四个进程全状态分片得到的 4 MB，报告为完整内存峰值。指出三个漏项。
4. 十进制工作权重从 1.00 开始，每次减去 0.003。比较连续两次立即保留两位小数，与保留更精细主副本的结果。
5. 一次恢复加载了权重和矩状态，却重新初始化数据采样器。还能声称什么？尚未建立哪项恢复证据？
6. 激活检查点让保存的激活减半，却让算术耗时增加 30%。在什么容量条件下仍然有用？为什么不能据此称每次迭代更快？

27.9 参考解答

1. 保存的激活加倍为 2,621,440 字节，加上不变的 16,000,000 字节模型状态，总计 18,621,440 字节。真实注意力和工作区开销需另建估算，这个答案只服从题目给定的假设。
2. 梯度总和为 $3\times2+1\times(-2)=4$，除以四个事件得 1。等权平均两个均值会得到 0，让单个事件与三个事件具有相同影响。
3. 例如本地保存的激活、临时收集的权重，以及通信／内核缓冲区。碎片和规划余量还可能增加占用。长期分片存储只是其中一项，并不是实测峰值。
4. 立即舍入时，两次更新后都还是 1.00。主副本则保留 0.997、0.994，第二次的工作副本变成 0.99。这个十进制类比说明了逐次丢掉小变化与先积累变化的差别。
5. 可以说从已保存的模型／优化器数值状态继续训练，但不能证明后续数据序列相同，也不能证明与连续运行一致。应保存采样器状态，并在声明的环境中比较后续更新。
6. 若原先运行超过可用内存，而减少激活后实际峰值进入容量范围，重算就可能让训练变得可行。更多算术可能使每次迭代更慢；放得下与跑得快是不同标准。

27.10 参考文献

• Micikevicius 等，Mixed Precision Training，2018。第 3.1—3.2 节说明 FP16 条件下主权重副本与损失缩放分别解决什么问题。
• Rajbhandari 等，ZeRO，2020。第 3、5 节区分重复的状态类别、分片和额外内存项。
• Narayanan 等，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，2021。第 2 节介绍并行维度和流水线调度；本章没有把论文的集群性能当作自己的测量。