09 · 显存优化全家桶

⏱ 28 分钟 ｜ ⭐ OOM 时按这个顺序试

🎯 一句话

OOM 不是一个问题，是四个问题 —— 参数、梯度、优化器状态、激活， 每一块的解法完全不同。 第 3 章教你算它们各占多少，这一章教你怎么把每一块压下去。

🩺 一、先诊断：你的 OOM 属于哪一类

import torch
torch.cuda.reset_peak_memory_stats()
train_step()
print(f"峰值 {torch.cuda.max_memory_allocated()/1e9:.1f} GB  "
      f"已保留 {torch.cuda.max_memory_reserved()/1e9:.1f} GB")

# ⭐ 更有用的：完整的显存快照，能看出每一块是谁占的
torch.cuda.memory._record_memory_history()
train_step()
torch.cuda.memory._dump_snapshot("mem.pickle")
# 拖到 https://pytorch.org/memory_viz 里看可视化

⚠️ allocated 和 reserved 的差距就是碎片。 差距超过 20% 说明碎片严重，见文末。

🧰 二、七件武器（按性价比排序）

① 梯度检查点 —— 最常用

# 🧩 骨架：`model` 来自你自己的代码，这一段只看写法
# 整个模型
from torch.utils.checkpoint import checkpoint_sequential
out = checkpoint_sequential(model.layers, segments=4, input=x)

# ⭐ HuggingFace 一行
model.gradient_checkpointing_enable()

# ⭐ 更精细：选择性检查点（PyTorch 2.x）
#    只重算便宜的（LayerNorm/激活函数），保留贵的（矩阵乘）
from torch.utils.checkpoint import create_selective_checkpoint_contexts

💡 选择性检查点值得一提：全量重算掉 30% 速度， 但只重算那些"算起来便宜、存起来贵"的操作， 常常能用 5–10% 的时间换到大部分的显存收益。

② ZeRO / FSDP —— 切分优化器状态

③ 激活值量化 / 更低精度

④ 换优化器

import bitsandbytes as bnb
opt = bnb.optim.AdamW8bit(model.parameters(), lr=3e-4)   # ⭐ 一行换掉

⑤ CPU Offload

⑥ 参数高效微调（PEFT）

🔑 这是显存问题最容易被忽略的解法： 很多人在硬啃全量微调的显存，而任务本身用 LoRA 就够了。

⑦ 减小 batch + 梯度累积

# 🧩 骨架：`loader` 来自你自己的代码，这一段只看写法
for i, batch in enumerate(loader):
    loss = model(batch) / accum_steps        # ⭐ 别忘了除
    loss.backward()
    if (i + 1) % accum_steps == 0:
        opt.step(); opt.zero_grad(set_to_none=True)

🧭 三、决策树（照着走）

🧩 四、显存碎片：那个"明明还有显存却 OOM"

# ⭐ 三个常用手段
import torch
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
#                                        ↑ PyTorch 2.x，对碎片非常有效

torch.cuda.empty_cache()          # 释放缓存（慢，别放在循环里）

# ⭐ 最有效的预防：让形状固定
#    变长序列 → 分桶（bucketing）到固定几个长度

⚠️ empty_cache() 不是解药：它会让分配器重新向驱动申请，反而变慢。 真正的解法是让分配模式变规律 —— 固定形状、分桶、expandable_segments。

📋 五、一张速查表

⭐ 注意 FlashAttention 和 BF16 这两行： 它们几乎没有代价 —— 应该是默认开启的，不算"优化手段"。 如果你还没开这两个就在纠结 offload，顺序错了。

🔗 和站内其他章的关系

✅ 检查点

1. 判断 OOM 属于哪一类的关键实验是什么？
2. allocated 和 reserved 的差距说明什么？
3. 梯度检查点的显存和时间各是什么量级？
4. 什么是选择性检查点？它好在哪？
5. 8-bit Adam 能省多少？代价是什么？
6. 为什么说 LoRA 是"最容易被忽略的解法"？
7. 为什么"减小 batch"是最后手段而不是第一手段？
8. 显存碎片的三个症状？最有效的预防是什么？
9. 速查表里哪两个手段"几乎无代价"？这意味着什么？

🛑 可以停在这里

⚡ 走神救援

⭐ OOM 不是一个问题，是四个（参数、梯度、优化器状态、激活），每块解法不同。

⭐⭐ 诊断只要一个实验：把 batch 调到 1，还 OOM 吗？ 还 OOM 就是参数或优化器的问题——调 batch 根本没用；不 OOM 就是激活的问题。另外 allocated 和 reserved 差得多就是碎片。

七件武器里最该记的四条：⭐ 梯度检查点把激活显存降一个量级、代价是多花两三成时间（选择性检查点只重算「算起来便宜、存起来贵」的那部分，性价比更高）；⭐ 切优化器状态（ZeRO/FSDP）按 stage 递进；⭐ 换 8-bit 优化器几乎是白拿——一行的事、质量损失极小；⭐⭐ LoRA 是最容易被忽略的解法——很多人在硬啃全量微调的显存，而任务用 LoRA 就够，可训练参数掉到百分之一量级，梯度和优化器状态跟着掉。

⚠️ 两个「最后手段」：CPU offload 要走 PCIe，慢一个量级，实在装不下才用；减小 micro batch 会让 warp 不够、通信占比上升，MFU 反而变差。

⭐ 决策树按那个诊断实验分叉：batch=1 还 OOM 走 LoRA → 8-bit 优化器 → 切分 → 并行 → offload；不 OOM 走梯度检查点 → FlashAttention → 缩短序列 → 序列并行。

碎片的症状很有辨识度：「空着好几 GB 却分配不了小一半的显存」、跑几小时后突然 OOM。⭐ 解法是让形状固定（分桶）加上换分配器策略；⚠️ empty_cache() 不是解药，反而更慢。

⭐ 最后一句是顺序问题：FlashAttention 和 BF16 几乎无代价，它们本来就该是默认开的——还没开就先纠结 offload，是把顺序搞反了。

下一节 👉 10-数据并行与AllReduce.md