按精度计算模型状态的理论内存占用

你好。上一课我们沿一次训练 step 追踪了参数如何参与前向、反向如何产生梯度，以及 DDP 与 FSDP 分别在哪里进行通信。现在需要把这条时间线变成一份可计算的模型状态内存账本：每个参数到底对应多少字节？为什么 BF16 训练并不意味着“每参数只占 2 字节”？又为什么 Adam 的状态往往比模型权重本身更昂贵？

本课只计算模型状态的理论载荷，即权重、梯度、主权重和优化器状态。激活、注意力临时张量、通信工作区、KV Cache、内存碎片和 FSDP 的短时参数物化会在后续课程中分别建模。这个边界很重要：状态账本给出“常驻负担”的下界，而不是一次训练 step 的完整显存峰值。

先确定：什么算模型状态，什么不算

令模型有 $N$ 个唯一的可训练参数元素。这里的“唯一”意味着共享或权重绑定的参数只能计数一次。例如，若 embedding 与 LM head 共享同一块存储，不能在 $N$ 中重复计算。

训练中的模型状态通常包括五类张量：

状态	记号	典型精度	为什么需要它
计算权重	$W_{model}$	BF16 或 FP16	前向和反向中的 GEMM、注意力等算子读取它
模型梯度	$G_{model}$	BF16、FP16 或 FP32	反向传播写入，并参与梯度累积或通信
主权重	$W_{master}$	常为 FP32	优化器实际更新的高精度权重副本
一阶矩	$m$	常为 FP32	Adam 对梯度均值的指数滑动估计
二阶矩	$v$	常为 FP32	Adam 对梯度平方的指数滑动估计

Adam 的状态来自其更新规则：

m_{t} = β_{1} m_{t - 1} + (1 - β_{1}) g_{t}

v_{t} = β_{2} v_{t - 1} + (1 - β_{2}) g_{t}^{2}

因此，对每一个参数元素，Adam 至少额外维护 $m$ 与 $v$ 两个同形状张量。若它们都以 FP32 保存，就是每参数额外 $8$ 字节。

与之相对，以下内容不是本课的模型状态账本：

前向保存的激活及其重计算元数据；
FlashAttention、归一化、路由和 GEMM 的临时工作区；
FSDP AllGather、ReduceScatter 或 All-to-All 的通信缓冲区；
CUDA/ASIC 运行时、内存分配器保留空间及碎片；
长上下文 rollout 中的 KV Cache。

后续做设备容量规划时，需要将这些部分逐项加到状态账本上；但如果连模型状态本身都没有算清楚，后续的容量模型就没有可靠基线。

为什么低精度权重仍常配 FP32 主权重

BF16 或 FP16 权重每元素只需 $2$ 字节，因而很适合前向与反向计算。但训练更新不是简单地将一个大数替换成另一个大数，而是不断累加很小的更新量：

W_{t + 1} = W_{t} - η \cdot g_{t}

当学习率与梯度的乘积很小，直接对低精度权重做加法可能因为尾数精度不足而完全不改变存储的权重值。FP16 的问题尤其明显；BF16 的指数范围较大，溢出和下溢风险较低，但其尾数也比 FP32 短，因此许多大规模训练实现仍保留 FP32 主权重。

典型更新路径是：

W_{model}^{BF16} \to G \to W_{master}^{FP32} \to cast (W_{model}^{BF16})

也就是说：

BF16 权重供前向和反向计算使用；
反向得到梯度；
优化器以 FP32 主权重、FP32 优化器状态执行更新；
更新后的主权重再转换回 BF16，作为下一轮计算权重。

这是一条数值稳定性与内存之间的折中路径。要注意：主权重并不是“同一个权重换了一个 dtype”的抽象概念；它通常是一块真实存在、与模型权重同形状的独立存储，因此必须单独计入内存。

Mixed Precision Training | Explanation and PyTorch Implementation from Scratch

观看 ExplainingAI 的《Mixed Precision Training | Explanation and PyTorch Implementation from Scratch》。这一段用一个具体的小更新说明了为何低精度权重不能总是直接更新，并展示了 FP32 主权重与低精度计算权重如何并存。

观看主权重机制。重点不是复现实现代码，而是确认两件事：低精度计算权重服务于前反向，而 FP32 副本服务于更新；因此二者在内存账本中必须分别计数。

通用公式：按“存活副本”而不是按模型名称计费

任何模型状态估算都可写为：

M_{state} = N (n_{W} b_{W} + n_{G} b_{G} + n_{M} b_{M} + n_{m} b_{m} + n_{v} b_{v})

其中：

$N$ ：唯一参数元素数；
$n_{x}$ ：该类状态在同一时刻实际存活的副本数；
$b_{x}$ ：每个元素的字节数；
BF16、FP16： $2$ 字节；
FP32： $4$ 字节。

最容易出错的地方是梯度。工程实现中可能同时有：

autograd 产生的低精度模型梯度；
用于梯度累积、归约或优化器更新的 FP32 主梯度缓冲区。

如果两者在峰值时刻同时存在，就必须都计入；如果实现直接将梯度写入或复用 FP32 缓冲区，则不能重复计数。“梯度是几字节”不是模型的固有属性，而是某个训练引擎在某个时刻的实际存储策略。

对使用 Adam 的常见配置，可以先建立如下理论表。

配置	权重	梯度	主权重	Adam $m, v$	合计
全 FP32 Adam	4	4	0	8	$16$ 字节/参数
BF16 权重、FP32 梯度、FP32 主权重、FP32 Adam	2	4	4	8	$18$ 字节/参数
FP16 权重与模型梯度，另有 FP32 主梯度、主权重和 Adam	2	2 + 4	4	8	$20$ 字节/参数

第二行是目标平台上非常值得作为基线的账本：BF16 用于模型计算，梯度和更新路径主要采用 FP32。

第三行体现了一个工程事实：即使模型权重、模型梯度都是 16 位，若训练系统又保留 FP32 主梯度缓冲区，状态负担不但不会是“2 加 2”，反而可能达到 $20$ 字节/参数。

这里的“理论内存”指张量内容的字节数，不含地址、对齐和分配器额外开销。容量规划时应将最终结果再留出运行时安全余量，但不能用模糊的安全余量代替清晰的状态分解。

5T MoE 的第一条硬约束：总参数决定状态，不是激活参数

你的目标模型有约 $5 \times 1 0^{12}$ 个总参数。即使 MoE 每个 token 只激活 Top-16 专家，所有专家的可训练权重、梯度和 Adam 状态仍必须存在于训练系统中。Top- $k$ 主要改变单 token 的计算量、路由通信和激活；它不会将未激活专家的训练状态从内存账本中删除。

先按十进制单位计算：

1 TB = 1 0^{12} bytes

若采用 BF16 计算权重、FP32 梯度、FP32 主权重和 FP32 Adam 状态：

B_{per param} = 2 + 4 + 4 + 4 + 4 = 18 bytes

那么全模型状态的逻辑总量为：

M_{state} = 5 \times 1 0^{12} \times 18 = 90 \times 1 0^{12} bytes = 90 TB

换算为二进制单位约为：

90 TB \approx 81.8 TiB

作为对照：

状态方案	字节/参数	5T 逻辑模型状态
全 FP32 Adam	$16$	$80 TB$
BF16 权重 + FP32 更新路径	$18$	$90 TB$
FP16 双梯度缓冲区实现	$20$	$100 TB$

这给出一个很有价值的架构结论：5T 模型不能仅靠“用 BF16”解决状态容量问题。 BF16 确实将计算权重从 FP32 的 $20 TB$ 降到 $10 TB$ ，但 Adam 的主权重、动量与方差依然构成数十 TB 的逻辑状态。状态分片不是锦上添花，而是必需条件。

DDP 中的“副本”与分布式优化器中的“分片”

在普通同步数据并行中，数据并行组大小设为 $d$ 。每个 rank 都保存完整模型状态：

M_{rank}^{DDP} = M_{state}

因此，增加 DDP rank 的数量会扩大总物理容量，却不会降低单卡模型状态内存。每张卡仍保留完整权重、完整梯度、完整主权重与完整 Adam 状态。

分布式优化器则首先分片主权重和优化器状态。Megatron Core 给出的一个常见 BF16 配置是：

B_{rank} = 6 + \frac{12}{d}

其中：

$6$ 字节/参数是每个 rank 上仍复制保存的 BF16 模型权重与 FP32 梯度；
$12/ d$ 字节/参数是按数据并行组分片的 FP32 主权重、FP32 一阶矩与 FP32 二阶矩。

Distributed Optimizer — Megatron Core

阅读 Megatron Core 的《Distributed Optimizer》。它给出了不同精度配置下每参数字节数的表，并把“常驻复制的部分”与“按数据并行度分片的部分”写成可直接使用的公式。

在页面开头的 “Distributed Optimizer” 部分，从内存公式与缓冲区说明读到参数与主梯度缓冲区的说明。重点核对表中的 6+12/d 形式，而不是只记忆数值。随后在 “Key Steps” 部分，阅读更新步骤，观察 ReduceScatter 后每个 rank 如何只更新自己拥有的那一段 FP32 状态，再通过 AllGather 让模型计算权重恢复一致。

若仅作说明性计算，取 $d = 64$ ，且暂时假定一个数据并行副本面对完整 $5$ T 参数：

B_{rank} = 6 + \frac{12}{64} = 6.1875 bytes/param

M_{rank} = 5 \times 1 0^{12} \times 6.1875 \approx 30.94 TB

相比未分片的 $90 TB$ ，分布式优化器显著减少了主权重与 Adam 状态；但 $30.94 TB$ 仍远超单设备容量，因为模型计算权重与主梯度仍被复制。

这说明“分片优化器状态”和“全分片模型状态”是不同层次的能力：

普通 DDP：权重、梯度、主权重、优化器状态均复制；
分布式优化器：通常先分片主权重和优化器状态；
FSDP/ZeRO-3 类方案：进一步分片参数和梯度，但会引入按需 AllGather、ReduceScatter 与临时物化峰值。

上一课提到的 FSDP 参数物化属于第三类；它的精确峰值模型需要同时考虑分片常驻状态和逐层通信缓冲区，后续分布式并行模块会专门处理。当前最重要的是先不把这三种状态布局混成一个公式。

如何读“参数与梯度缓冲区分片”图

图中将单个 rank 可见的参数与梯度缓冲区划成四个数据并行分片；绿色、黄色、蓝色和红色分别表示不同参数张量，图下方列出每个 DP rank 负责的全局索引、局部缓冲区索引和参数内索引。

这张图有三个容易被忽略的含义。

第一，分片边界不必与参数张量边界对齐。例如某个分片可能从 Param 1 的中间开始，到 Param 2 的中间结束。因此真实实现通常采用连续扁平缓冲区，而不是“每个完整层或完整张量恰好分给一个 rank”。

第二，图中的三类索引分别服务于不同职责：

world index：逻辑全局扁平缓冲区中的位置；
local index：当前 rank 所负责分片在本地缓冲区中的位置；
param index：该元素在原始参数张量内部的位置。

这也是检查点重分片、优化器状态迁移、BF16 到 FP8 权重发布时必须保存布局元数据的原因：只知道“参数名”不够，还必须知道它在分片缓冲区中的偏移和长度。

第三，不要因为图中出现了参数分块，就直接断言“每个 rank 永远只保留自己的参数块”。在一些分布式优化器实现中，模型计算权重缓冲区仍可能在各 rank 上完整存在，只是主权重与优化器更新的所有权按分片划分。判断内存时，应回到实现的真实生命周期：

哪些张量是完整副本？
哪些张量按哪个并行组分片？
哪些张量仅在前向、反向或更新的短暂阶段存在？
同一时刻是否有模型梯度与主梯度两份副本？

用于架构评审的一页状态账本

对后续的 8 卡、64 卡和更大规模方案，建议固定使用下面这张账本，而不是只写“模型为 BF16”。

字段	应填写的内容
参数计数 $N$	总参数、共享参数去重方式、是否包含全部 MoE 专家
计算权重	dtype、每参数字节数、是否复制或分片
反向梯度	dtype、是否有梯度累积缓冲区、是否有独立主梯度
主权重	dtype、是否存在、分片组大小
优化器状态	优化器类型、 $m$ 、 $v$ 的 dtype 与分片策略
常驻状态内存	按每 rank 计算的理论值
短时状态峰值	参数物化、通信工作区、转换缓冲区，单列估计
未计入项	激活、临时张量、KV Cache、运行时余量及其独立模型

对于目标 ASIC，状态账本还应增加一列“存储位置”：HBM、片上 SRAM、Host 内存或远端存储。因为“状态能否放得下”与“状态能否以足够带宽参与更新”是两个不同问题。尤其在高算力、低带宽平台上，将 FP32 主权重或 Adam 状态放到较慢存储，可能节省 HBM，却会把优化器更新变成新的吞吐瓶颈。

小结

本课建立了模型状态的基础内存公式：

M_{state} = N \times 每参数所有同时存活状态的字节数

关键结论是：

模型状态至少要分别清点计算权重、梯度、FP32 主权重与优化器状态。
BF16 或 FP16 权重并不意味着训练状态只需 $2$ 字节/参数；Adam 的 FP32 主权重、动量和方差常占主要部分。
对 BF16 权重、FP32 梯度、FP32 主权重和 FP32 Adam 状态，理论基线为 $18$ 字节/参数。
$5$ T 参数在该基线下对应约 $90 TB$ 的逻辑模型状态；MoE 的 Top- $k$ 不会消除未激活专家的状态存储。
DDP 扩大设备数量但不降低单 rank 状态；分布式优化器和 FSDP 的价值在于减少这些状态的复制，但其通信与短时峰值需要单独建模。
任何内存讨论都应明确：这是常驻状态、瞬时峰值，还是包含运行时余量的设备容量预算。

下一课将进入 MoE 的核心量纲区分：总参数量、每 token 激活参数量与实际计算量。这会解释为什么一个 5T 总参数模型的状态容量极大，但其单 token FLOPs 又可能接近远小于 5T 的稠密模型。

Recommended resources