显存优化：4bit / gradient checkpointing 增强版

Unsloth 在单卡 24 GB 上能训 Llama-3-70B QLoRA，关键就是把"显存占用从哪来"这件事算到了像素级。 本章拆解一份训练任务的显存预算，把 Unsloth 的每个省显存招都讲明白。

6.1训练显存账分解

LoRA 训练时一张卡上的显存可以分成 4 块：

$$ M_\text{total} = M_\text{weights} + M_\text{optimizer} + M_\text{gradients} + M_\text{activations} $$

对 LoRA 训练，"权重" 和 "激活/logits" 是两大头。Unsloth 对这两块都有专门优化。

6.2第一招：4bit 量化模型权重

原理

4bit 量化把权重从 fp16 (2 bytes) 压到 nf4 (0.5 bytes)，显存降 4 倍。bitsandbytes 的 nf4 数学：

$$ W_\text{quant}[i] = \mathrm{round}\Big(\frac{W[i]}{s}\Big),\quad s = \frac{\max|W_\text{block}|}{7} $$

每 64 个权重共享一个 scale $s$（block-wise）。反量化时乘回 scale 恢复 fp16。

Unsloth 怎么用

model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/llama-3-8b-bnb-4bit",   # ★ 直接用预量化版
    load_in_4bit = True,              # 或 from base 实时量化
)

用 unsloth/*-bnb-4bit 预量化模型的好处：

• 跳过实时量化（节省几分钟）；
• 量化质量经过 Unsloth 验证（有些原版 nf4 有 bug，Unsloth 修过）；
• 下载量小 4 倍。

显存收益

6.3第二招：8bit Adam（adamw_8bit）

Adam 优化器对每个可训练参数维护 $(m, v)$ 两份状态，fp32 时占 $8\,\Psi$ bytes。 bitsandbytes 的 adamw_8bit 把 $m, v$ 压到 8bit，优化器状态显存降 4 倍。

args = TrainingArguments(
    ...,
    optim = "adamw_8bit",     # ★ Unsloth 推荐默认
)

对 LoRA 训练效果不太明显（LoRA 可训参数本来就少），但对全参 / 长序列是必开。

6.4第三招：unsloth gradient checkpointing

这是 Unsloth 的"独门暗器"。原版 PyTorch gradient_checkpointing=True 的逻辑：

• 前向丢弃中间激活（每个 transformer block 重算）；
• 反向时重算前向恢复激活；
• 激活显存降到 $\sqrt{L}$ 量级；
• 训练时间 +25-35%。

Unsloth 的 "unsloth" 模式更激进：

model = FastLanguageModel.get_peft_model(
    model,
    use_gradient_checkpointing = "unsloth",   # 不是 True，是字符串 "unsloth"
)

额外做的事：

1. 把每层的输入激活 offload 到 CPU 而不是丢掉；
2. 反向时按需 reload 到 GPU；
3. 激活显存额外降 30%；
4. 训练时间仅多 5-10%（offload 通信被 overlap 掉）。

对长序列（≥ 4K）训练效果尤其明显。

6.5第四招：chunked cross-entropy（再讲一遍）

第 5 章讲过 chunked CE 的实现。从显存角度看：

$$ M_\text{logits} = B \cdot S \cdot V \cdot 4\ \text{bytes}\ (\text{fp32}) $$

$B$=2, $S$=2048, $V$=128000 时 = 2 GB；fp32 反向时还要存一份梯度，翻倍到 4 GB。 chunked CE 把这一项降到接近 0。

在 Unsloth 里默认开启（patch 后），用户无需配置。

6.6第五招：低秩 LoRA + 选择性 target_modules

LoRA 可训参数量：

$$ \#_\text{trainable} = N_\text{layers} \cdot \sum_\text{target} 2\,d \cdot r $$

典型推荐：

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],   # 全 Linear
    lora_alpha = 16,
)

6.7第六招：缩小 max_seq_length

显存里激活的部分约 $O(B \cdot S \cdot H)$。把 max_seq_length 砍掉一半，激活显存约降一半。

但有个陷阱：max_seq_length 应该设到训练数据实际需要的长度，不要无脑设大。例：

• alpaca / sharegpt 数据：512-1024 够用；
• 长指令 / 长对话：2048；
• 长文档处理：4096 - 32768。

看数据集 token 长度分布：

from datasets import load_dataset
import numpy as np

ds = load_dataset("yahma/alpaca-cleaned", split="train")
lengths = [len(tokenizer.encode(x["instruction"]+x["output"])) for x in ds.select(range(1000))]
print(f"max={max(lengths)}, p95={np.percentile(lengths, 95)}, p99={np.percentile(lengths, 99)}")

把 max_seq_length 设到 p95 / p99 附近即可。

6.8第七招：缩小 per_device_train_batch_size

显存最直接的开关。如果 OOM，降到 1，配合 gradient_accumulation 不损失等效 batch：

args = TrainingArguments(
    per_device_train_batch_size = 1,            # 单卡 batch 1
    gradient_accumulation_steps = 8,             # 累 8 步反传一次
    # 等效 batch = 1 × 8 = 8
)

6.9显存对照实例

Llama-3-8B QLoRA + bs=2 + seq=2048 + r=16 的实测峰值（4090 24G）：

用同样 4090 多出来的 14 GB 你可以：

• 把 batch_size 升到 8；
• 或把 seq_length 升到 8K；
• 或把 r 升到 64；
• 或把模型换成 13B；
• 或加上推理 vLLM engine（在线 RL）。

6.10显存自查流程

训练前用以下脚本估算：

import torch

def estimate_memory(num_params_b, lora_r, max_seq, batch, dtype_bytes=2):
    weights = num_params_b * 1e9 * dtype_bytes / 1e9        # 4bit 时 ×0.25
    lora_params = num_params_b * 1e9 * 0.005                 # LoRA r=16 ≈ 0.5%
    optim = lora_params * 8                                  # Adam fp32
    grads = lora_params * 2                                   # bf16
    activations = batch * max_seq * 4096 * 32 * 2 / 1e9       # 极简公式
    print(f"weights={weights:.1f}GB, optim={optim/1e9:.2f}GB, grads={grads/1e9:.2f}GB, act={activations:.1f}GB")

estimate_memory(num_params_b=8, lora_r=16, max_seq=2048, batch=2)

训练时实测：

print(f"Peak: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

6.11显存不够时的优先级排序

1. 开 4bit 量化（最大头）；
2. 开 use_gradient_checkpointing="unsloth"；
3. 用 adamw_8bit 优化器；
4. 降 batch size 到 1，加 gradient_accumulation；
5. 降 max_seq_length 到 p95；
6. 降 LoRA r 到 8；
7. 减少 target_modules（去掉 down_proj 等大矩阵）；
8. 真不行 → 换更小模型 / 多卡 Unsloth Pro。

6.12这章带走的

• 4bit 量化 = 权重显存 ÷4，是 QLoRA 的基础；
• adamw_8bit = 优化器状态显存 ÷4；
• use_gradient_checkpointing="unsloth" 比 True 再省 30%；
• chunked CE 自动开，大词表必备；
• OOM 时按"4bit → ckpt → 8bit Adam → bs → seq → r → target_modules" 顺序砍；
• 同硬件 Unsloth 比 HF + PEFT 省 30-50% 显存。