从0训练LLM · 现代大模型体系 阅读 + 实验 ≈ 40 分钟 页内实验 ×4
C9 精度与显存:一个参数要几字节 第 5 课

激活为什么比参数还占显存?

激活比参数还大
前向的中间结果要存下来给反向用,每层 s·b·h·(34 + 5as/h) 字节:和参数量无关,由序列长度、batch、宽度、层数决定,注意力分数那一项还带 s²。锚点一条 4,096 token 的序列,按标准块公式近似就是 162.5 GiB,注意力分数占七成多。FlashAttention 不写出分数表,重计算用多算一遍前向换不存。
❓

本课一问:前四课把参数、梯度、矩阵乘都压下来了。可一旦把训练序列拉长,显存还是会爆——这时候占地方的已经不是参数了。激活为什么比参数还占显存?

这一课先不碰:把激活切到多张卡上(C10 的序列并行、上下文并行);线性注意力为什么没有 s² 这一项(C11);推理时的 KV cache(C3、C13)。

从已有出发

前四课一直在算「每个参数几个字节」。可训练一步留在显存里的,还有前向时存下来、留给反向用的中间结果——激活。点下面的「前向」:

从已有出发 · 训练一步在显存里留下什么(BF16 混合精度) 本课点任一环节看它在 Qwen/Qwen3.8-27B 上占多少
精度与显存的演进:每砍一次位宽,都在一个真数字上失败一次
  1. 起点 FP32 全精度 + AdamW 深度学习早期的默认做法
    解决了范围到 ±3.4 × 10³⁸,1 附近的间距 2⁻²³ ≈ 1.2 × 10⁻⁷,数值上几乎不用操心。
    卡在每个参数 4 + 4 + 4 + 4 = 16 字节(权重、梯度、Adam 的 m、v):锚点 27,781,427,952 个参数就是 444.5 GB,是 80 GB 单卡的 5.6 倍,还没算激活。
  2. 第一代 FP16 混合精度:FP32 主权重 + loss scaling Micikevicius et al. 2017
    解决了矩阵乘和激活用 16 位,带宽与激活显存减半;FP32 主权重接住小更新,loss 乘 S 把小梯度推出下溢区,更新前再除回来。
    卡在指数只有 5 位:fp16(10⁻⁸) = 0,fp16(1 + 10⁻⁴) = 1,fp16(70,000) = inf;缩放倍数要一直调,溢出就得跳过这一步。
  3. 第二代 BF16:用尾数换指数 Google Brain 2019 · Kalamkar et al. 2019
    解决了8 位指数,范围与 FP32 相同:不再下溢,不需要 loss scaling。2026 年的默认训练精度。
    卡在1 附近间距粗到 1/128:bf16(257) = 256,bf16(1 + 10⁻³) = 1——主权重和优化器状态仍要 FP32,账本还是 16 字节/参数。
  4. 第三代 FP8 + 细粒度块缩放 Micikevicius et al. 2022(格式)· DeepSeek-V3 2024(超大规模训练)
    解决了激活每 1 × 128、权重每 128 × 128 一个缩放:张量 [1000, 0.01, 0.02, 0.5] 里,0.01 的误差从整张量缩放的 12.8% 降到 0.45%;DeepSeek-V3 相对 BF16 的 loss 相对误差 < 0.25%(官方自报)。
    卡在只有 256 个位型,E4M3 最大 448;块的形状要跟着离群值走——激活梯度改按 128 × 128 块量化,16B MoE 训到约 300B token 发散(附录 B.2);主权重和梯度仍是 FP32。
  5. 岔路 · 2025–2026 4 位:MX / NVFP4 块缩放 + QAT OCP MX 2023 · NVFP4 2025 · Kimi K3 / DeepSeek-V4 2026
    解决了MXFP4 每 32 个数共用一个 2 的幂(4.25 位/参数),NVFP4 每 16 个数一个 E4M3 缩放;K3、DeepSeek-V4 在后训练阶段对 MoE 专家权重做 QAT。
    卡在E2M1 只有 16 个码点:块 [0.3, −1.2, 2.7, 0.05] 用 MX 缩放后,0.05 → 0;全 FP4 预训练仍是前沿。
  6. 另一半账 激活:重计算与 FlashAttention Chen et al. 2016 · Korthikanti et al. 2022 · Dao et al. 2022 你在这里
    解决了不把 s × s 的注意力分数存下来:锚点一条 4,096 token 的序列,激活从 162.5 GiB 降到 42.5 GiB(选择性重计算 / FlashAttention);全量重计算每层只存输入,只剩 2.5 GiB(标准块公式近似)。
    卡在重计算用多算一遍前向换显存;可 444.5 GB 的模型状态一张 80 GB 卡仍然放不下 → C10 并行。
数字口径:舍入结果在构建时按 IEEE 754 / OCP 规则现算(先存 FP32,再就近舍入、平局取偶);显存由 anchor-model.json 的 param_total 现算,1 GB = 10⁹ 字节;激活按 Korthikanti et al. 2022 式 1(标准块近似,s = 4,096、b = 1);DeepSeek-V3 的数字出自其技术报告 §3.3 与附录 B(本地 PDF,官方自报)。

贯穿例句回来了。用「小明把苹果给了小王」这 5 个 token 训练锚点模型一步,按本课要讲的公式,64 层的激活一共约 55.9 MB——和上一课 444.5 GB 的模型状态比,可以忽略。可训练不会只喂 5 个 token:把这句话换成一条 4,096 个 token 的文档,激活就变成 162.5 GiB。s 从 5 变到 4,096,放大了约 820 倍,激活却放大了约 3,121 倍。多出来的那部分,就是本课要找的 s²。

先自己设计:前向的中间结果全部存下来

反向传播要用到前向的中间结果:算一个矩阵乘的权重梯度,要用到它的输入;算 softmax 的梯度,要用到它的输出。最直接的做法:前向每算出一个中间结果,就存着,等反向用完再释放。

  1. ① 朴素方案
  2. ② 它在哪失败
  3. ③ 改进方案
朴素方案前向的中间结果全部存下来
全存:模型状态 + 激活

这不只是锚点的账。ZeRO 论文算过一个更早的例子:GPT-2 1.5B(序列 1K、batch 32)的激活约 60 GB,而它的模型状态只有 24 GB 。R4 调研用 Qwen3-8B 的维度(s = 4,096、b = 1)算出 109 GiB、注意力分数占 82.5%,组件里切到 Qwen3-8B 可以复核。

这个公式从哪来

Korthikanti 等人(NVIDIA,2022)把一个标准 Transformer 块前向时要存的每一个张量逐个数了一遍。s 是序列长度,b 是 batch,h 是宽度,a 是注意力头数,每个数按 16 位(2 字节)存:

每层激活(字节)· Korthikanti et al. 2022 式 1
· ( + )
👆 公式里带下划线的每一项都可以点——点哪项，讲哪项。
前提写在论文 §4:16 位存储、dropout 掩码每个 1 字节、MLP 宽 4h、标准多头注意力,忽略 LayerNorm 的均值方差这类小缓冲。锚点是 GQA、SwiGLU、没有 dropout、还有 48 层线性注意力——把它的 h、a、L 代进这个公式,只是近似。

把「注意力分数」这一项画成张量,就看得见 s² 长在哪:

一个全注意力层里的分数表(锚点 Qwen3.8 的真实维度,b = 1)
O = softmax(Q · Kᵀ ÷ √dh) · V
输入输出中间结果 渐变 = 分数(连续值)
边长按对数比例缩放,块边上的数字是真实值;同一维度符号的边在全图等长。
维度含义 真实规模 小规模
a Q 头数(num_attention_heads) 24 2
s 序列长度(一条 4,096 token 的训练序列) 4096 4
dh 每头维度(head_dim) 256 3
P 一共 24 × 4,096² = 402,653,184 个数:按 BF16 存 805.3 MB;按公式的 5 字节(含 dropout 的两份)是 2.0 GB——这只是一层、一条序列。输出 O 只有 25,165,824 个数。
🎯

回到上面的堆叠条,试三件事:① 选锚点、「全存」,把 s 从 5 拖到 131,072,看「注意力分数」那一段怎么从看不见变成压倒一切;② 固定 s = 4,096,把 batch 从 1 拖到 8:激活正比增长,而模型状态那几段纹丝不动;③ 勾选「只有 16 个全注意力层带 s² 项」:锚点 64 层里只有 16 层是全注意力,激活从 162.5 GiB 降到 72.5 GiB——另外 48 层是线性注意力,没有 s × s 的表(C11 揭晓;这里 48 层仍按 34sbh 近似)。

⚠

常见误解一(旧版课程的错误):「激活和参数量成正比。」

旧版的显存组件把激活估成「每参数 3 字节」,于是拖参数量滑杆时,各部分比例永远不变。实际上激活由 s、b、h、L 决定:同一个模型,序列从 1K 拉到 8K,激活可以翻几十倍,参数一个没变;而注意力分数那一项和参数量完全无关,只和头数、序列长度有关。

常见误解二:「batch 翻倍,显存就翻倍。」

只有激活随 b 翻倍;16 字节的模型状态与 batch 无关。所以大模型训练里,「一张卡能放多大的 batch」是先扣掉模型状态,再看剩下的空间够放几条序列的激活。

改进一:梯度检查点——只存一部分,反向时重算

2016 年 Chen 等人的做法:前向时只在少数「检查点」上存激活,其余的算完就扔;反向走到哪一段,就从最近的检查点把那一段的前向再算一遍。n 层的网络选约 √n 个检查点,激活显存从 O(n) 降到 O(√n),代价是多一次前向。原文的例子:1,000 层的 ResNet 从 48G 降到 7G,多花约 30% 的时间 。

大模型里最常见的是每一层都做检查点:每层只存进入这一层的输入(16 位,每层 2sbh),反向时逐层重算。锚点 s = 4,096 时,激活从 162.5 GiB 降到 2.5 GiB,反向时再临时多一层完整激活。代价是整个前向多算一遍:C6 第 1 课说训练算力约 6ND(前向 2ND + 反向 4ND),开了全量重计算就是约 8ND——那里留的「重计算」的疑问在这里兑现。Korthikanti 的论文实测,全量重计算让训练慢 30% 到 40%。

改进二:选择性重计算——只重算最占地方、最不费算力的那一块

看公式的两项:34sbh 那部分来自矩阵乘的输入,重算它们要把矩阵乘再做一遍,很贵;5as²b 那部分来自 softmax 和 dropout,占的显存多,算起来却便宜。Korthikanti 等人因此只重算注意力的核心(QKᵀ、softmax、dropout、乘 V),其余照存:每层激活从 sbh(34 + 5as/h) 降到 34sbh。锚点 s = 4,096:42.5 GiB。论文把它和序列并行(C10)一起用,激活降到原来的约 1/5,重计算的额外开销减少 90% 以上;530B 模型在 2,240 张卡上的 MFU 从全量重计算的 42.1% 提到 54.2% 。

改进三:FlashAttention——分数表根本不写进显存

2022 年 Dao 等人从另一个方向解决同一件事:注意力分数表 S、P 在朴素实现里要写进显存(HBM),再读回来做 softmax、乘 V,反向再读一遍。FlashAttention 把 Q、K、V 分块搬进芯片上的小而快的 SRAM,在块内算分数、用「在线 softmax」逐块累积结果,s × s 的表从不写进显存;反向时需要的分数,用前向存下的 softmax 统计量现算回来 。

它是精确的,不是近似;显存从 O(s²) 变成 O(s);而且因为省掉了大量显存读写,反而更快——GPT-2(序列 1K)端到端快 3 倍 。C0 第 4 课说过,N² 这笔账后来被两条路分别处理,其中一条就是「不把整张表写出来」——在这里兑现;但 N² 次乘加一次都没省,另一条路(让大多数层根本没有这张表)在 C11。

在上面的堆叠条里,「FlashAttention」与「选择性重计算」在存储上是同一个效果(都去掉了 5as²b 那一项);差别在算力:前者少了显存读写、更快,后者要把分数重算一遍。

⚠

常见误解三:「FlashAttention 是一种近似注意力。」

不是。它改的是「怎么算」,不是「算什么」:分块 + 在线 softmax 与标准注意力在数学上等价,论文标题里就写着 exact attention。输出与标准实现的差别只来自浮点加法顺序的不同(和 C2 第 1 课那个「差值恰为 0」的注脚是同一件事)。

常见误解四:「重计算是免费的。」

全量重计算多一次前向,训练慢 30% 到 40%;选择性重计算便宜得多,但也不是零。显存和算力在这里是交换关系:省下的每一个字节,都要用重新算一遍来换。

一次迭代里,显存怎么涨、怎么落

上面的堆叠条是某一刻的截面。把时间轴展开:前向时激活逐层攒起来,反向时逐层释放、同时梯度逐层长出来,最后 optimizer.step()。这是按公式推出来的推演,不是实测:

一次迭代里,显存怎么涨、怎么落(推演)

🎯

试三件事:① 点「播放一次迭代」,看峰值出现在哪一刻(前向刚结束:激活攒到最多);② 把「第 1 次迭代」换成「第 2 次」:起点高了 222.3 GB——Adam 的 m、v 在第一次 step() 才分配,从第二次迭代起一直在(PyTorch 的显存博客记录过同一现象);③ 把「激活」换成「全量重计算」,看峰值降了多少、峰值的位置挪到了哪。

我们自己的模型上,实验脚本分阶段读了 4 次真实显存(建模型后 / 前向后 / 反向后 / 第一步后)。数据生成后,切到「我们的小模型」,黑点就是实测值,可以和推演的曲线对照:

项目实测 · 4 个显存截面

证据三问

问本课的回答
① 它解决什么问题?长序列训练时激活放不下:不存注意力分数表,或只存检查点、反向时重算
② 数学上为什么可能有效?激活公式是逐张量记账 ;FlashAttention 与标准注意力数学等价
③ 实验上为什么被选中?选择性重计算 + 序列并行把 530B 的 MFU 从 42.1% 提到 54.2%;FlashAttention 在 GPT-2 上快 3 倍;两者今天是训练的标配

一句话

📌

激活由 s、b、h、L 决定,和参数量无关;注意力分数那一项带 s²,序列一长就压倒一切——FlashAttention 不写出这张表,重计算用多算一遍前向换不存。

🚦

先停在这里:记住「激活 ∝ s·b·h·L,另有 a·s²·b 的分数表;FlashAttention 去掉它,重计算用算力换显存」就可以去 C10。
再深一层:去实验课,把我们自己模型的显存截面和公式对一遍账。

下一个缺口

激活压下去之后,27.8B 的 444.5 GB 模型状态仍放不进一张 80 GB 的卡——一张卡不够,怎么办? → C10 并行:一张卡不够怎么办

第 1 课那本 16 字节的账,在 C10 第 2 课(ZeRO)被切到多张卡上;本课的激活,在 C10 的序列并行、上下文并行里被沿序列切开。

你可能会问