在训练开始之前，你必须知道一件事：这个模型训练得下去吗？显存够不够？
四大显存消费者
训练时，GPU 显存被四样东西占据：
1. 模型参数
Sophia 约 10 亿个参数。混合精度下需要 FP32 主副本 + BF16 工作副本：
FP32: 10 亿 × 4 字节 = 4.0 GB
BF16: 10 亿 × 2 字节 = 2.0 GB
合计：6.0 GB
2. 梯度
反向传播产生的梯度，每个参数一个，BF16 存储：
10 亿 × 2 字节 = 2.0 GB
3. 优化器状态
AdamW 对每个参数维护 m 和 v 两个状态，FP32 存储（需要高精度）：
10 亿 × 2 × 4 字节 = 8.0 GB
Muon 的额外状态也在这个量级（具体取决于实现）。
4. 激活值（中间计算结果）
前向传播时，每一层的输入和中间结果需要存下来给反向传播用。这是最不确定的一项——取决于 batch size 和序列长度。
粗略估算：一个 token 经过一层 Transformer 大约产生 6,000 到 10,000 个 BF16 数字的激活值。5 万 token × 28 层 × 8,000 ≈ 11.2 亿个 BF16 数字 ≈ 2.2 GB。
但实际中可以用激活检查点（activation checkpointing）技术来节省：不存所有层的激活，只存每隔几层的。需要时重新计算——用计算时间换显存空间。
总显存估算
参数： 6.0 GB
梯度： 2.0 GB
优化器： 8.0 GB
激活值： 2-4 GB（取决于是否用激活检查点）
PyTorch 开销：1-2 GB
─────────────────
总计： 19-22 GB
RTX 5090 有 32 GB。留有约 10 GB 的余量——这个余量让训练有呼吸空间，不会因为偶尔的显存碎片化或意外的尖峰而崩溃。
推理时的显存
推理时不需要梯度、不需要优化器状态、不需要存激活值给反向传播。只需要：
模型参数（BF16）：2.0 GB
KV cache：取决于上下文长度
这就是为什么推理比训练轻松得多——显存需求只有训练的十分之一。
下一章节来看怎么衡量 GPU 的利用效率——你的代码真正用到了 GPU 多少算力？