训练显存、混合精度与 LoRA

本章对应课件：《3.1 大模型工程专题》后半部分。核心知识点全部按统一模板展开：直觉、公式、逐项拆解、shape、数字算例、交互演示、代码、误区、面试问题。

17.1 FP32、FP16、BF16：数值格式 ★

数字算例（7B 模型的权重显存）

逐项拆解（位分布）

矩阵 shape

格式不影响张量形状，只影响每个元素的字节数：

17.2 为什么 BF16 很重要

FP16：精度较高，但 exponent（指数）范围相对有限。

BF16：mantissa（尾数）精度低一些，但 exponent 范围更接近 FP32。

直觉对比：

• FP16 像「刻度很细但尺子很短」——精确但容易超出量程；
• BF16 像「刻度粗但尺子很长」——范围大，适合梯度这种动态范围极大的数值。

现在很多 GPU（A100、H100、H800）上的大模型训练大量使用 BF16。

17.3 Mixed Precision：混合精度训练 ★

经典 mixed precision 模型（FP16 时代）

这是理解 mixed precision 的经典教学模型（以早期 FP16 训练为代表）：

Forward            → FP16
Backward           → FP16（梯度也是 FP16）
Master weights     → FP32（额外保存一份高精度参数）
Optimizer states   → FP32
Loss scaling       → GradScaler 放大 loss，防止 FP16 梯度下溢

为什么需要 master weights？ 假设学习率 $\eta = 10^{-5}$，某参数 $w = 0.5$，更新量：

FP16 在 0.5 附近的精度（最小间隔）约为 $2^{-11} \approx 4.9 \times 10^{-4}$——比更新量大一个数量级。如果直接在 FP16 上原地更新，$w + \Delta w$ 会被舍入回 $w$，模型「学不动」。所以在 FP32 副本上累加更新，再转成 FP16 供前向使用。

数字算例小结

16 / 18 bytes per parameter 都只是「明确假设下的估算」——见 17.4 节的口径说明框。

# 类型：【Skeleton】需要 GPU
scaler = torch.amp.GradScaler("cuda")      # FP16 必须：放大 loss，防止梯度下溢

for x, y in loader:
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(x)
        loss = loss_fn(logits, y)
    scaler.scale(loss).backward()          # 放大梯度
    scaler.step(optimizer)                 # 先 unscale，再 step（内部处理 inf/nan）
    scaler.update()
    optimizer.zero_grad()

现代实践：BF16 / torch.amp

BF16 的指数范围与 FP32 相同，通常不需要 GradScaler：

# 类型：【Skeleton】需要 GPU
for x, y in loader:
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(x)
        loss = loss_fn(logits, y)
    loss.backward()                        # 直接 backward，无需 scaler
    optimizer.step()
    optimizer.zero_grad()

17.4 训练显存估算 ★

公式（每参数）

逐项拆解

数字算例

7B 训练就超过单卡 80GB——所以大模型训练必须用分布式技术：

17.5 两个显存开关：梯度累积与激活检查点

开关一：梯度累积（Gradient Accumulation）

小显存装不下大 batch 时：把一个大 batch 拆成 $k$ 个 micro-batch，梯度累加 $k$ 次后再更新一次参数。

# 类型：【Runnable】核心逻辑（完整训练循环见第 13 章 Lab 9）
accum = 8
for micro_step, (x, y) in enumerate(loader):
    loss = criterion(model(x), y) / accum        # 缩放：保证累加后与平均一致
    loss.backward()                              # 梯度累加进 .grad
    if (micro_step + 1) % accum == 0:            # 攒够 accum 个才更新
        clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

开关二：激活检查点（Activation Checkpointing / 梯度检查点）

正常：      Forward → 保存全部 activation → Backward 直接使用
Checkpoint：Forward → 只保存边界 activation → 丢弃中间结果
           Backward → 重新执行对应区间的 forward → 得到中间值 → 继续反传

核心 trade-off：更少的激活显存 ↔ 更多的重计算（时间）。粒度越粗（检查点区间的层数越多），保存得越少、重算得越多。

# 类型：【Skeleton】需要模型与输入
from torch.utils.checkpoint import checkpoint
h = checkpoint(block, x, use_reentrant=False)   # 该 block 的中间激活不再全部保存

FP8 简报：H100 起支持 FP8（E4M3/E5M2）训练，进一步减半显存与带宽需求，但需要缩放与兼容的 kernel 支持。它属于「前沿选项」——本节不展开，知道它存在即可。

17.6 LoRA：低秩微调 ★

公式

其中：

逐项拆解

矩阵 shape

数字算例（小矩阵，看得见的低秩）

设 d = k = 4，r = 1：

观察：ΔW 的每一行都是 A 的倍数（行之间线性相关）——这就是「低秩」：无论怎么训练，$BA$ 只能表示少数方向的组合。

17.7 为什么 LoRA 省参数 ★

公式与算例

假设 $W \in \mathbb R^{4096 \times 4096}$：

全量微调：

LoRA（r = 8）：

占比：

关键：省的不只是权重显存，更是优化器状态（Adam 的 m/v 是参数量的 8 字节/参数）——LoRA 只需为 A/B 保存优化器状态。

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=8,                      # 秩
    lora_alpha=16,            # 缩放系数（实际用 α/r · BA）
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
model = get_peft_model(base_model, config)
model.print_trainable_parameters()
# trainable params: 4,194,304 || all params: 6,742,609,920 || trainable%: 0.0622

17.8 LoRA 为什么可行

它背后的假设：

模型微调所需要的参数变化，并不一定占满整个高维参数空间，可能只需要在一个较低维的方向上调整。

所以 $\Delta W$ 近似是 low rank。这是 LoRA 最核心的思想。

矩阵秩的直觉

实践细节：

• 初始化：$A$ 用高斯随机，$B$ 初始化为 0 → 训练开始时 $\Delta W = 0$，不破坏原模型；
• 缩放：实际使用 $\frac{\alpha}{r} BA$，$\alpha$ 控制更新幅度；
• 作用位置：通常加在 Attention 的 $W_Q, W_K, W_V, W_O$ 上；
• r 的选择：8～64 常用，任务越复杂/数据越多可以越大。

17.9 本章总结