09 · 显存优化全家桶

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 里看可视化

   ⭐ 一个关键的判断：把 batch 调到 1，还 OOM 吗？

   还 OOM  → 是【参数/梯度/优化器状态】的问题
             → 调 batch 没用，要用 ZeRO / 量化 / 换优化器
   不 OOM  → 是【激活】的问题
             → 梯度检查点、序列并行都能救

   正常：前向时保存【每一层】的激活，反向时用

   检查点：只保存【少数几个】节点的激活
          反向时从最近的节点【重新前向计算】出中间激活

   ⭐ 显存：O(L) → O(√L)
   ⚠️ 时间：+20~35%（多一次前向）

# 整个模型
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

   [第 3 章]算过：优化器状态是最大的一块（12 字节/参数）

   ZeRO 把它切开分给 N 张卡：
   Stage 1：切优化器状态      → 省 4x
   Stage 2：+ 切梯度          → 省 8x
   Stage 3：+ 切参数          → 省 Nx  ⭐

   🔗 第 11 章详讲

   BF16 激活 → FP8 激活：显存再减半
   ⚠️ 需要 H100 且要处理 scaling（第 6 章）

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

   把优化器状态 / 参数放到 CPU 内存，需要时搬过来

   ⚠️ 走 PCIe（第 3 章：比 HBM 慢约 40 倍）
   → 只在【实在装不下】时用，会显著拖慢

   💡 例外：推理时 offload 不常用的专家（MoE）是划算的

   ⭐ 如果你只是【微调】而不是预训练，这是最省的路：

   LoRA：冻结原权重，只训练低秩增量 A·B
   → 可训练参数降到 0.1~1%
   → 梯度和优化器状态跟着降到 1%  ⭐⭐

   QLoRA：base 模型用 4-bit 量化 + LoRA
   → 单张 24GB 卡能微调 33B 模型

   micro_batch=1，累积 32 步 → 等效 batch=32

   ⚠️ 但要知道代价：
   · micro batch 太小 → GPU 利用率下降（第 2 章：warp 不够）
   · 通信次数不变但计算变少 → 通信占比上升

   ⭐ 所以这是【最后手段】，不是第一手段

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 了
   │
   ├─ batch=1 还 OOM？
   │   ├─ 是 → 参数/优化器的问题
   │   │      ① 是微调吗？→ 【LoRA / QLoRA】⭐ 最省事
   │   │      ② 换 8-bit Adam / Adafactor
   │   │      ③ 上 ZeRO-2 → ZeRO-3（第 11 章）
   │   │      ④ 还不行 → 张量并行 / 流水线并行（第 12-13 章）
   │   │      ⑤ 最后 → CPU offload（慢）
   │   │
   │   └─ 否 → 激活的问题
   │          ① 【梯度检查点】⭐ 先试这个
   │          ② FlashAttention（第 8 章，把 O(n²) 降到 O(n)）
   │          ③ 序列长度能不能减？能不能 packing？
   │          ④ 序列并行（超长序列时）
   │          ⑤ 减小 micro batch + 梯度累积
   │
   └─ 时好时坏 / 显存够但报 OOM？ → 【碎片】，见下

   症状：
   · 报错说"tried to allocate 2GB，free 8GB"但还是失败
   · 训练跑了几小时后突然 OOM
   · reserved 比 allocated 大很多（>20%）

   原因：PyTorch 的缓存分配器把显存切成了很多小块，
        没有【连续】的大块可用

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

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

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