03. 训练显存账本：8B 模型一张卡训得动吗

上一篇走完了一个训练 step，途中出现了好几份和参数一样大的数据：BF16 权重、梯度、FP32 主权重、AdamW 的两份动量，还有为反向保存的激活。推理系列的显存账本算过，Llama-3.1-8B 推理时权重只占约 15 GiB，一张 H100 绰绰有余。这一篇算训练的账：同一个模型全参数微调要多少显存，为什么一张 80 GB 的卡放不下，LoRA 和 QLoRA 各省掉了哪几笔，激活重计算又拿什么换了显存。读完以后，你能拿到一个模型的参数量和训练配置，估算出它需要几张什么样的卡。

训练比推理多出三笔账

推理时显存里主要是权重和 KV Cache。训练时没有 KV Cache（一次前向就处理整条序列），但多出了三类东西，一共四个桶：

• 权重：前向和反向用的那份参数，混合精度下是 BF16。
• 梯度：反向给每个参数算出的梯度，和权重一样多。
• 优化器状态：AdamW 的 FP32 主权重、一阶动量 m、二阶动量 v，每个参数三份。
• 激活：前向时保存、留给反向用的中间结果，和参数量无关，随 batch、序列长度和层数增长。

此外还有框架开销：CUDA 运行环境、通信缓冲区、PyTorch 缓存分配器留着没还的块，以及显存碎片。

前三个桶只和参数量有关，叫模型状态，可以精确算出来；激活和开销取决于训练配置，只能估算后实测。

每个参数 16 字节

按上一篇的混合精度 AdamW 逐项列出每个参数要存什么：

内容	格式	每参数字节数	为什么要存
权重	BF16	2	前向、反向计算
梯度	BF16	2	交给优化器
主权重	FP32	4	小改动在 BF16 里会被舍掉
一阶动量 m	FP32	4	梯度方向的滑动平均
二阶动量 v	FP32	4	梯度大小的滑动平均
合计		16

这个 16 字节的口径来自 ZeRO 论文（2019 年，微软 DeepSpeed 团队），当时用 FP16，今天换成 BF16 字节数不变。代入 Llama-3.1-8B 的 8,030,261,248 个参数：

模型状态 = 参数量 × 16 字节
        = 8.03×10⁹ × 16 ≈ 128.5 GB ≈ 119.7 GiB
其中 BF16 权重 16.1 GB，梯度 16.1 GB，FP32 主权重与两份动量各 32.1 GB

一张 H100 标称 80 GB，系统实际可用约 79.6 GiB。还没开始算激活，模型状态就超出了 40 GiB。推理只要 15 GiB 权重，训练要 8 倍，这是“能推理的卡不一定能训练”的根本原因。

真正的大头是优化器那 12 字节，占了 3/4。这也解释了后面两条省显存的路：要么让大部分参数不需要优化器状态（LoRA），要么把这些状态分到多张卡上（ZeRO/FSDP）。

激活：随 batch 和序列长度增长

上一篇说过，反向要用到前向的中间结果。每一层的注意力和前馈网络里，矩阵乘法的输入、归一化前后的值、激活函数的输入，都要留到反向时再用。这些数的个数正比于“这一层处理了多少个 token × 每个 token 的向量宽度”。

Megatron 团队 2022 年的论文给过一个 GPT 类结构的估算式，每层激活约：

每层激活字节数 ≈ s × b × h × (34 + 5 × a × s ÷ h)
s：序列长度   b：micro-batch 大小   h：隐藏维度   a：注意力头数

34 × s × b × h 来自层内各处保存的张量（16 位、每个 token 若干个 h 宽的向量）；5as/h 那一项是注意力分数矩阵，随序列长度平方增长。现在的训练几乎都用 FlashAttention，它不保存这个矩阵，反向时现算，这一项就可以去掉。Llama 的前馈层结构和 GPT 不同，系数 34 只能当量级参考。代入 8B、序列 2048、micro-batch 1：

34 × 2048 × 1 × 4096 × 32 层 ≈ 9.1 GB（估算）

还有一笔容易漏掉的：输出层的 logits。每个位置对 128,256 个 token 各有一个分数，算交叉熵时常转成 FP32：

2048 × 128,256 × 4 字节 ≈ 1.05 GB（每条 2048 token 的序列）

词表越大这笔越重。Llama 3 的词表约是 Llama 2（32,000）的 4 倍，长序列训练时 logits 能占到可观的比例。

激活的关键在于它怎样缩放：和 micro-batch 成正比，和序列长度成正比（不用 FlashAttention 时还有平方项），和层数、隐藏维度成正比。micro-batch 从 1 改成 4，激活从 9 GB 变成 36 GB；序列从 2048 改成 8192，再乘 4。前面算的模型状态是固定成本，激活是随配置变化的成本，OOM（显存不足）报错多半出在这里。

激活重计算：用时间换显存

如果每层都不存中间结果，只存这一层的输入，反向走到某一层时，先用存下的输入把这一层的前向重新算一遍，得到中间结果，再求梯度，用完就丢。这叫激活重计算，也叫激活检查点或梯度检查点（gradient checkpointing）：

显存从“所有层的激活”降到“所有层的输入加一层的激活”。每层输入是 s × b × h 个 BF16 数：

各层输入：2 字节 × 2048 × 4096 × 32 层 ≈ 0.54 GB
重算中的一层：34 × 2048 × 4096 ≈ 0.29 GB
合计约 0.8 GB，原来约 9.1 GB（都是估算，未计 logits）

代价是多做一次前向。反向的计算量大约是前向的 2 倍（既要算对激活的梯度，也要算对权重的梯度），原本一步是 1 + 2 = 3 份，现在是 1 + 1 + 2 = 4 份，多出约 1/3。ZeRO 论文引用的数字是 33%，Megatron 论文实测全量重计算带来 30%—40% 的时间开销。实际框架还可以只重算一部分开销大、计算便宜的操作（选择性重计算），在两者之间取折中。

LoRA：大部分参数不需要优化器

LoRA 冻结原模型的全部参数，在选定的线性层旁边加一对小矩阵，只训练这对小矩阵（原理和超参数在第 10 篇）。资料助手用秩 r = 16，作用于每层全部 7 种线性层（注意力的 q、k、v、o 和前馈的 gate、up、down）。一个形状为“输入 × 输出”的线性层，新增 (输入 + 输出) × r 个参数：

q、o：(4096 + 4096) × 16 = 131,072，各一个
k、v：(4096 + 1024) × 16 = 81,920，各一个（8 个 KV 头 × 128 维 = 1024）
gate、up、down：(4096 + 14336) × 16 = 294,912，各一个
每层：131,072 × 2 + 81,920 × 2 + 294,912 × 3 = 1,310,720
32 层：约 4194 万，占 80.3 亿的 0.52%

只有这 4194 万个参数需要梯度和 AdamW 状态，按每参数 16 字节算约 0.67 GB。冻结的底座只需要 BF16 权重这一份，16.1 GB，没有梯度、主权重和动量。

模型状态从 128.5 GB 降到约 16.8 GB，但激活没怎么变少。LoRA 分支插在每一层里，要给第 1 层的 LoRA 求梯度，反向就得穿过全部 32 层，每层的激活仍要保存或重算。QLoRA 论文举过 7B 模型的例子：LoRA 参数只占 26 MB，batch 为 1 时 LoRA 反向需要的输入梯度却有 567 MB。所以 LoRA 训练时激活重计算照样重要。

QLoRA：把冻结的底座压到 4 bit

既然底座冻结不更新，它就只需要“能参与前向计算”，不需要高精度地接受小改动。QLoRA（2023 年）把底座的线性层量化成 4 bit 存储，计算时按块反量化成 BF16，LoRA 分支仍是 BF16。

它用的格式叫 NF4（4-bit NormalFloat），16 个取值按正态分布的分位点排布，适合大致呈正态分布的权重。每 64 个权重为一块，共用一个缩放常数。如果常数用 FP32 存，每个权重要分摊 32 ÷ 64 = 0.5 bit。QLoRA 再对这些常数做一次 8 bit 量化（每 256 个常数一块，块内再共用一个 FP32 常数），叫双重量化：

每权重分摊的常数开销 = 8 ÷ 64 + 32 ÷ (64 × 256) ≈ 0.127 bit
每个线性层权重 ≈ 4 + 0.127 = 4.127 bit

bitsandbytes 库的 4 bit 层默认就是块大小 64 并开启双重量化。词嵌入不是线性层，不会被量化；输出层 lm_head 通常也保留 BF16。代入 8B：

32 层线性层：6,979,321,856 × 4.127 ÷ 8 ≈ 3.60 GB
词嵌入 + 输出层：2 × 128,256 × 4096 × 2 字节 ≈ 2.10 GB
底座合计 ≈ 5.7 GB ≈ 5.3 GiB（不开双重量化约 6.0 GB）

Llama 3 的词表很大，两个不量化的词表矩阵占了底座的三分之一以上；词表只有 32,000 的老模型，这一项要小得多。QLoRA 还用了分页优化器：显存不够时把优化器状态临时换到 CPU 内存，防止重计算带来的显存峰值直接 OOM。

省显存不是免费的。4 bit 反量化让每步计算变慢；训练出的 LoRA 是对着量化后的底座学的，部署时如果换回 BF16 底座或合并进权重，行为可能有细微差别，需要用评测集确认（第 10 篇）。

三种做法放在一起

把四个桶按比例画出来，横轴单位是 GiB，两条虚线是 24 GB 消费级显卡和 H100 的容量：

全参数微调那一条里，黄色的 FP32 主权重和两份动量占了大半，它们全部来自优化器。LoRA 把这三块和梯度几乎删光，剩下的蓝色底座成了主角；QLoRA 再把底座压到三分之一。对照两张卡的容量（序列 2048、micro-batch 1 的估算，LoRA 和 QLoRA 开启重计算）：

做法	模型状态	激活与开销	合计（估算）	1 张 H100（79.6 GiB）	24 GB 消费级卡
全参数 + AdamW	119.7 GiB	约 11 GiB	约 131 GiB	放不下，开重计算也放不下	放不下
LoRA，BF16 底座	15.6 GiB	约 3.3 GiB	约 19 GiB	宽裕，可加大 batch 或序列	能跑，余量只有几 GiB
QLoRA，NF4 底座	5.9 GiB	约 3.3 GiB	约 9 GiB	很宽裕	能跑，序列可到数千

表里的激活和开销是估算，具体数字随框架版本、注意力实现、是否把 logits 分块计算而变化，误差可以达到几个 GiB。正式训练前，用目标配置里最长的样本跑几步，读 torch.cuda.max_memory_allocated() 或 nvidia-smi 的峰值，才是可靠的数字。

一个常见的 OOM：最长的那条样本

资料助手团队在 24 GB 卡上跑 LoRA，前 200 步一切正常，第 213 步报 OOM。检查发现，这一步的 micro-batch 里恰好有一份把整本员工手册贴进去的样本，长 8192 token，而且数据加载时把同一个 micro-batch 的 4 条样本都补齐（padding）到最长那条的长度。

按前面的缩放关系算一下。平时 4 条样本各几百 token，补齐后通常不超过 1024，一个 micro-batch 约 4 × 1024 = 4096 个位置；这一步是 4 × 8192 = 32,768 个位置，多了 8 倍。前面估算的 LoRA 激活（开重计算：各层输入、重算中的一层，加上 logits）约 1.85 GB 即 1.7 GiB，对应 2048 个位置，每个位置约 0.85 MiB：

平时：4096 个位置 × 约 0.85 MiB ≈ 3.4 GiB
第 213 步：32,768 个位置 × 约 0.85 MiB ≈ 27 GiB（估算）

模型状态没变，激活一下子超过了整张卡。前 200 步没出事，只是因为长样本还没被抽到。

判断这类 OOM 的方法是看它是否出现在固定的某些 step、而且和样本长度相关：记录每步的最大序列长度和显存峰值，两条曲线会对上。处理办法按代价从小到大排：截断或拆分超长样本；按长度分桶或做样本打包，减少 padding；减小 micro-batch、增大梯度累积；开启激活重计算；最后才是换更大的卡。估算显存时要按最长样本算，而不是平均长度。

放不下时：把状态分到多张卡

全参数微调 8B 在单卡上放不下，但这些状态不必挤在一张卡上。普通的数据并行让每张卡各存一份完整的 16 字节状态，加卡并不能解决放不下的问题。ZeRO（DeepSpeed）和 PyTorch 的 FSDP 把优化器状态、梯度乃至权重切成 N 份，每张卡只存自己那份，需要时再通过通信取回：

8 张卡完全切分：128.5 GB ÷ 8 ≈ 16.1 GB 每卡（模型状态）
再加上每张卡自己的激活和开销

这样 8 张 H100 就能全参数微调 8B，代价是卡间通信。具体怎样切、通信多少，是第 08 篇的内容。在单卡上还有一些折中办法，比如把优化器状态存成 8 bit、或者卸载到 CPU 内存，它们用精度或速度换显存，适合资源确实受限的场景。

训练显存分四个桶：BF16 权重和梯度各 2 字节，AdamW 的 FP32 主权重和两份动量共 12 字节，合计每参数 16 字节，8B 模型仅这一项就是 128.5 GB；激活随 micro-batch、序列长度、层数和隐藏维度增长，激活重计算多花约 1/3 的计算把它压到原来的一小部分。LoRA 让绝大多数参数不再需要梯度和优化器状态，QLoRA 再把冻结的底座压到 4 bit，8B 模型因此能在一张 24 GB 的卡上训练。显存决定了能不能训，下一篇换一个问题：训完要多少算力、多少卡、多少天，以及为什么今天的小模型都读了远超“最优”数量的数据：04. 算力账与规模定律。