面试-训练显存占用估计

5. 训练显存占用

极简公式：

训练总显存 =
模型权重显存
+ 梯度显存
+ 优化器状态显存
+ 前向激活值显存
+ 临时 buffer / CUDA runtime / 显存碎片

和推理最大的区别是：训练不仅要存模型，还要存梯度、优化器状态，并且激活值不能立刻释放，因为反向传播要用。

---

5.1 全参数训练显存

一句话： 全参数训练 = 权重 + 梯度 + 优化器状态 + 激活值。

以 BF16 训练、AdamW 优化器为例：

1）模型权重

模型权重 = 参数量 × 2 byte

Qwen3-1.7B：

1.7B × 2 byte ≈ 3.4GB

2）梯度显存

每个可训练参数都要保存梯度。

梯度 = 参数量 × 2 byte

所以：

1.7B × 2 byte ≈ 3.4GB

3）AdamW 优化器状态

AdamW 通常为每个参数维护：

m 一阶动量
v 二阶动量

很多训练框架中 optimizer state 默认是 FP32：

AdamW状态 ≈ 参数量 × 2 × 4 byte

所以：

1.7B × 2 × 4 byte ≈ 13.6GB

如果还保留 FP32 master weight，则还要额外：

1.7B × 4 byte ≈ 6.8GB

4）激活值显存

训练时激活值是大头之一。

推理时：

with torch.no_grad():
    out = model(x)

中间激活可以尽快释放。

训练时不能这样，因为反向传播需要：

前向每层保存激活值 → 反向时用这些激活值计算梯度

激活值大致和下面因素正相关：

batch_size × seq_len × hidden_size × 层数 × dtype字节数

并且 Attention 中间矩阵还会受到：

batch_size × num_heads × seq_len × seq_len

影响。

所以训练时 seq_len 一旦变长，显存上涨非常明显。

---

5.2 全参数训练示例：Qwen3-1.7B

假设：

BF16 全参数训练
AdamW
batch=1
seq_len=2048
不开 ZeRO
不开 gradient checkpointing

估算：

权重:        3.4GB
梯度:        3.4GB
Adam状态:    13.6GB
激活值:      4GB ~ 8GB
CUDA开销:    0.5GB ~ 1GB

总计：

约 25GB ~ 30GB+

如果保留 FP32 master weight：

再加约 6.8GB
总计可能到 32GB ~ 38GB+

所以 2B 级模型推理可能 6GB 左右能跑，但全参数训练通常需要 24GB 以上，长序列或大 batch 会继续上涨。

---

6. LoRA / QLoRA 训练显存

6.1 LoRA 显存为什么小

LoRA 的核心思想：

冻结原模型参数，只训练少量低秩矩阵 A 和 B

原来的线性层：

y = W x

LoRA 后：

y = W x + B A x

其中：

W 冻结，不更新
A、B 可训练

所以 LoRA 不需要为全模型参数保存梯度和 Adam 状态，只需要为 LoRA 参数保存。

---

6.2 LoRA 训练显存公式

LoRA训练总显存 =
基础模型权重显存
+ LoRA参数显存
+ LoRA梯度显存
+ LoRA优化器状态显存
+ 激活值显存
+ runtime开销

重点区别：

全参数训练: 梯度和优化器状态覆盖整个模型
LoRA训练: 梯度和优化器状态只覆盖 LoRA 小参数

LoRA 参数量通常只占原模型的：

0.1% ~ 2%

所以优化器状态显存会小很多。

---

6.3 LoRA 示例：Qwen3-1.7B

假设：

基础模型 BF16
LoRA rank=8/16/32
batch=1
seq_len=2048

估算：

基础模型权重:      3.4GB
LoRA参数+梯度:     几十 MB ~ 几百 MB
LoRA优化器状态:    几十 MB ~ 几百 MB
激活值:            4GB ~ 8GB
CUDA开销:          0.5GB ~ 1GB

总计大约：

8GB ~ 13GB

注意：LoRA 虽然大幅减少「梯度 + 优化器状态」，但不会自动消灭训练激活值 。如果 seq_len 很长，激活值仍然会成为大头。

---

6.4 QLoRA

QLoRA 进一步把基础模型权重量化到 4bit：

基础模型权重: BF16 → INT4/NF4

但 LoRA 参数通常仍用 BF16/FP16 训练。

QLoRA训练总显存 =
4bit基础模型权重
+ LoRA参数/梯度/优化器状态
+ 激活值
+ 量化反量化buffer
+ runtime开销

Qwen3-1.7B：

BF16权重 ≈ 3.4GB
4bit权重 ≈ 0.85GB

所以 QLoRA 可能进一步降到：

6GB ~ 10GB

但具体取决于 seq_len、batch、是否开 gradient checkpointing、视觉模型输入分辨率等。

---

7. Gradient Checkpointing

训练时激活值太大，可以用 gradient checkpointing。

普通训练：

前向保存大部分激活值 → 反向直接用

Gradient checkpointing：

前向少保存激活值 → 反向时重新计算一部分前向

优点：

显存下降明显

缺点：

训练速度变慢，通常慢 20% ~ 40%

它本质上是：

用更多计算换更少显存

---

8. 预训练和微调的区别

8.1 预训练

目标： 让模型从大规模通用数据中学习语言、世界知识、视觉语言对齐能力。

数据形式通常是：

海量文本
图文对
网页数据
代码数据
视频/多模态数据

训练目标通常是 next token prediction：

给定前面的 token，预测下一个 token

loss 是标准交叉熵：

Loss = -log P(真实下一个token | 前文)

例如文本：

输入: 今天天气很
目标: 好

模型需要提高：

P("好" | "今天天气很")

预训练特点：

数据量极大
训练 token 数极大
通常全参数训练
成本最高
决定模型基础能力

---

8.2 微调 / SFT

SFT：Supervised Fine-Tuning，有监督微调。

目标不是从零学语言，而是让模型学会某种任务格式、回答风格、领域能力。

你的猫咪视觉识别数据就是典型 SFT：

{
  "role": "user",
  "content": "请识别图中猫咪状态..."
}
{
  "role": "assistant",
  "content": "{\"cats_visible\": 1, ...}"
}

SFT 的 loss 通常只算 assistant 部分：

user prompt 部分不算 loss
assistant answer 部分算 loss

也就是：

输入: user prompt + assistant answer
Loss mask:
user prompt         → -100，不参与loss
assistant answer    → 正常token id，参与loss

交叉熵形式仍然是：

Loss = -Σ log P(assistant真实token | image + prompt + 前文token)

对于你的 JSON SFT，本质是让模型学习：

图像 + 指令 → 合法 JSON 结构 + 正确字段值

---

8.3 预训练 vs 微调

项目	预训练	微调 / SFT
目标	学基础能力	学任务格式和领域行为
数据	海量通用数据	高质量任务数据
参数	通常全参数	全参数 / LoRA / QLoRA
成本	极高	较低
loss	next token loss	assistant answer token loss
作用	决定模型上限	改善特定任务表现
例子	学语言、视觉知识	学会输出猫咪 JSON schema

---

9. 训练显存汇总对比

以 Qwen3-1.7B、BF16、seq_len=2048 为例：

方式	权重	梯度	优化器状态	激活值	总显存估计
推理 no_grad + FA2	3.4GB	0	0	0.3~0.6GB	约 6.5GB
全参数训练	3.4GB	3.4GB	13.6GB+	4~8GB	25~38GB+
LoRA	3.4GB	很小	很小	4~8GB	8~13GB
QLoRA	约0.85GB	很小	很小	4~8GB	6~10GB

一句话总结：

推理显存主要看：权重 + KV Cache
训练显存主要看：权重 + 梯度 + 优化器状态 + 激活值
全参训练贵在梯度和Adam状态
LoRA省在只训练少量参数
QLoRA进一步省在基础模型权重量化
seq_len越长，激活值越容易成为训练显存大头