""" ============================================================================= 第 2 讲 · 代码 5：显存核算 —— 参数、梯度、优化器状态、激活值 ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 7 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 训练时显存里要放四样东西： 1. 参数（parameters） ：模型本身 2. 梯度（gradients） ：每个参数一个梯度，大小和参数一样 3. 优化器状态（optimizer state）：Adam 为每个参数额外存两个数（一阶矩、二阶矩） 4. 激活值（activations） ：前向传播的中间结果，反向传播时要用 前三项只和参数量有关；第 4 项还和批次大小、序列长度有关。 这个文件： - 算一算"每个参数要几个字节"，以及 8 块 H100 最多能训练多大的模型； - 实测 PyTorch 的 AdamW 到底存了多少状态； - 用 PyTorch 的"保存张量钩子"精确测量反向传播需要保存多少激活值； - 验证两种省显存的技巧：梯度累积（结果完全等价）、激活检查点（用计算换显存）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 05_memory_accounting.py """ import math import torch from torch.utils.checkpoint import checkpoint from torch.utils.flop_counter import FlopCounterMode def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def count_parameters(model): """ 统计模型的参数个数。 【PyTorch 小知识】 model.parameters() 逐个给出模型里所有需要训练的参数张量（nn.Parameter）。 """ return sum(p.numel() for p in model.parameters()) def make_mlp(num_layers, dim): """ 构造一个简单的深层网络：num_layers 个 [线性层 → ReLU]。 【参数】 num_layers : 层数（>= 1） dim : 每层宽度 【返回】 nn.Sequential """ layers = [] for _ in range(num_layers): layers += [torch.nn.Linear(dim, dim, bias=False), torch.nn.ReLU()] return torch.nn.Sequential(*layers) # ============================================================================= # 第 1 部分：每个参数要几个字节？ # ============================================================================= # 【思路】 # 使用 Adam 类优化器训练时，每个参数至少要： # 参数本身 + 梯度 + Adam 的一阶矩 m + Adam 的二阶矩 v # 它们各自用什么精度，取决于训练方案： TRAINING_SETUPS = { # 名字: (参数, 梯度, FP32 主权重, Adam m+v) 各自的字节数 "全 FP32": (4, 4, 0, 8), "BF16 参数和梯度 + FP32 Adam 状态": (2, 2, 0, 8), "混合精度（BF16 副本 + FP32 主权重）": (2, 2, 4, 8), } def demo_bytes_per_param(): """演示 1：不同训练方案下每个参数的字节数，以及 8 块 H100 最多能训练多大的模型。""" print_header("演示 1：每个参数要占几个字节？") total_mem = 8 * 80e9 # 8 块 H100，每块 80 GB print(f" {'方案':<34}{'参数':>4}{'梯度':>4}{'主权重':>6}{'Adam':>6}{'合计':>6} 8×H100 能放下的最大模型") for name, parts in TRAINING_SETUPS.items(): per_param = sum(parts) max_params = total_mem / per_param print(f" {name:<34}{parts[0]:>5}{parts[1]:>5}{parts[2]:>7}{parts[3]:>7}{per_param:>7} 约 {max_params / 1e9:.0f}B 参数") print(""" 注意：这还**没算激活值**，也没算框架本身和临时缓冲区占用的显存。实际能训练的模型要更小。 - 最常见的"混合精度"方案每个参数要 16 字节：前向和反向用 BF16 副本计算（快、省带宽）， 但优化器用 FP32 主权重来更新（避免小更新被舍入掉，见 02 号文件演示 5b）。 - 第 7 讲的 ZeRO / FSDP 会把这 16 字节切分到多块 GPU 上，从而训练大得多的模型。""") # ============================================================================= # 第 2 部分：实测 AdamW 的状态 # ============================================================================= def optimizer_state_bytes(optimizer): """ 统计优化器状态占用的字节数。 【思路】 PyTorch 优化器把每个参数的状态存在 optimizer.state[参数] 这个字典里， 例如 AdamW 存 "exp_avg"（一阶矩 m）、"exp_avg_sq"（二阶矩 v）和 "step"（步数）。 遍历这些状态里的张量，把字节数加起来。 【边界】必须先调用过至少一次 optimizer.step()，状态才会被创建（"懒创建"）。 """ total = 0 for state in optimizer.state.values(): for v in state.values(): if torch.is_tensor(v): total += v.numel() * v.element_size() return total def demo_adam_state(): """演示 2：AdamW 的状态是参数的 2 倍。""" print_header("演示 2：实测 AdamW 存了多少状态") model = make_mlp(4, 256) opt = torch.optim.AdamW(model.parameters(), lr=1e-3) n = count_parameters(model) print(f" 参数：{n:,} 个 × 4 字节 = {n * 4:,} 字节") print(f" step() 之前优化器状态：{optimizer_state_bytes(opt):,} 字节（还没创建）") model(torch.randn(8, 256)).sum().backward() opt.step() s = optimizer_state_bytes(opt) print(f" step() 之后优化器状态：{s:,} 字节 = 参数字节数的 {s / (n * 4):.2f} 倍") print(" （多出的一点点是每个参数张量一个 'step' 计数器）") for k, v in next(iter(opt.state.values())).items(): print(f" 状态 '{k}'：形状 {tuple(v.shape)}，dtype {v.dtype}") # ============================================================================= # 第 3 部分：精确测量激活值 # ============================================================================= def measure_saved_activations(fn, params): """ 测量执行 fn() 时，autograd 为反向传播保存了多少字节的**激活值**。 【思路】 PyTorch 在前向传播时，会把反向传播需要的张量"存起来"。 torch.autograd.graph.saved_tensors_hooks 允许我们在每次"存"的时候插入一个回调函数（pack_hook）， 这样就能看到到底存了哪些张量。 - 同一个张量可能被多个运算保存（比如 ReLU 的输出又是下一个线性层的输入），按内存地址去重； - 参数本身也会被保存（计算输入梯度要用 W），但参数不算"激活值"，排除掉。 【参数】 fn : 无参数函数，执行一次前向传播，返回输出 params : 模型参数的列表（用于排除） 【返回】 (字节数, 输出) 【边界】 这测量的是"前向结束时为反向保存的张量"，不包括反向传播过程中临时产生的梯度张量， 也不包括框架的内存池开销。它反映的是激活值显存的主要部分。 """ param_ptrs = {p.data_ptr() for p in params} seen = {} def pack_hook(t): ptr = t.data_ptr() if ptr not in param_ptrs: seen[ptr] = t.numel() * t.element_size() return t def unpack_hook(t): return t with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook): out = fn() return sum(seen.values()), out def demo_activation_memory(): """演示 3：激活值随批次大小和层数线性增长。""" print_header("演示 3：激活值 —— 与批次大小、层数成正比") dim = 512 print(f" 网络：L 层 × [Linear({dim},{dim}) → ReLU]，float32") print(f" {'L':>4}{'批次 B':>8}{'实测激活(字节)':>18}{'(L+1)×B×D×4':>16}{'参数(字节)':>14}") for L, B in [(4, 32), (8, 32), (8, 64), (16, 64)]: model = make_mlp(L, dim) x = torch.randn(B, dim) act_bytes, _ = measure_saved_activations(lambda: model(x), list(model.parameters())) formula = (L + 1) * B * dim * 4 print(f" {L:>4}{B:>8}{act_bytes:>18,}{formula:>16,}{count_parameters(model) * 4:>14,}") print(""" 每一层保存一个 B×D 的张量（ReLU 的输出，它同时是下一个线性层的输入），外加最开始的输入 x， 所以是 (L+1)×B×D 个数。 - 层数翻倍 → 激活翻倍；批次翻倍 → 激活翻倍。 - 参数量与批次无关。批次足够大时，激活值会超过参数、成为显存的大头。 【大模型联系】Transformer 每层保存的激活远不止一个 B×D（还有注意力分数、MLP 中间层等）， 并且与序列长度有关，第 3 讲会给出公式。""") # ============================================================================= # 第 4 部分：梯度累积 # ============================================================================= def demo_gradient_accumulation(): """演示 4：梯度累积 —— 用 4 个小批次模拟 1 个大批次，梯度完全一样，激活只要 1/4。""" print_header("演示 4：梯度累积") torch.manual_seed(0) model = make_mlp(4, 128) params = list(model.parameters()) x = torch.randn(64, 128) y = torch.randn(64, 128) def loss_fn(xb, yb): return ((model(xb) - yb) ** 2).mean() # 方式 A：一个大批次（64） act_big, loss = measure_saved_activations(lambda: loss_fn(x, y), params) loss.backward() grads_big = [p.grad.clone() for p in params] model.zero_grad() # 方式 B：4 个小批次（每个 16），每个的损失除以 4，梯度会自动累加到 .grad 里 # 【PyTorch 小知识】多次调用 backward() 时，梯度会**累加**到 .grad 中，而不是覆盖。 # 这既是常见 bug 的来源（忘了 zero_grad），也是梯度累积的实现基础。 n_micro = 4 act_small = 0 for xb, yb in zip(x.chunk(n_micro), y.chunk(n_micro)): a, loss = measure_saved_activations(lambda: loss_fn(xb, yb) / n_micro, params) act_small = max(act_small, a) # 同一时刻只需要保存一个小批次的激活 loss.backward() grads_acc = [p.grad.clone() for p in params] max_diff = max((g1 - g2).abs().max().item() for g1, g2 in zip(grads_big, grads_acc)) print(f" 大批次（64）的激活：{act_big:,} 字节") print(f" 小批次（16）的激活：{act_small:,} 字节（约 1/{act_big / act_small:.0f}）") print(f" 两种方式梯度的最大差异：{max_diff:.2e} ← 只是浮点舍入误差，数学上完全相等") print(""" 【边界】损失要除以小批次个数（因为 mean 是在每个小批次内部求的）。忘了除，梯度就会大 4 倍。 代价：4 次小的前向/反向比 1 次大的慢（矩阵变小，硬件利用率下降，见 04 号文件演示 4）。""") # ============================================================================= # 第 5 部分：激活检查点 # ============================================================================= def checkpoint_memory_units(num_layers, segment): """ 激活检查点的显存估算（单位：一层的激活大小）。 【思路】 每 segment 层设一个"检查点"：前向时只保存每段的输入（共 L/segment 份），段内的激活全部丢掉。 反向时，逐段重新做一遍前向（重计算），临时保留这一段内的 segment 份激活，用完即释放。 峰值显存 ≈ L/segment（检查点）+ segment（正在反向的那一段）。 这个和在 segment = √L 时最小，约为 2√L。 【参数】 num_layers : 层数 L segment : 每段的层数，1 <= segment <= L 【返回】 float，峰值显存（以"一层激活"为单位） """ return num_layers / segment + segment def demo_activation_checkpointing(): """演示 5：激活检查点 —— 用重计算换显存。""" print_header("演示 5：激活检查点（activation checkpointing）") L = 16 print(f" L = {L} 层时，不同分段大小的峰值激活显存（单位：一层的激活）：") print(f" 不做检查点：{L}") for seg in [1, 2, 4, 8, 16]: print(f" 每 {seg:>2} 层一个检查点：{checkpoint_memory_units(L, seg):>5.1f}") print(f" 最优分段 ≈ √L = {math.isqrt(L)}，峰值 ≈ 2√L = {2 * math.isqrt(L)}；" f"L = 100 时，从 100 降到约 {2 * 10}") print(""" 【边界】这里假设"一层内部的激活"和"这一层的输入"一样大（正是本文件这个简单网络的情况）， 所以"每层一个检查点"（17）反而比不做（16）还多。真实的 Transformer 层内部要保存注意力、 MLP 中间层等很多张量，内部激活是输入的十几到几十倍，那时"每层只存输入"能省下大部分显存。""") # 用 FLOP 计数器验证：检查点会多做一次前向 torch.manual_seed(0) dim, B = 256, 32 blocks = torch.nn.ModuleList([make_mlp(1, dim) for _ in range(L)]) x = torch.randn(B, dim) def run(use_ckpt): h = x for blk in blocks: # 【PyTorch 小知识】checkpoint(函数, 输入) 执行前向时不保存函数内部的激活， # 反向时再重新执行一遍这个函数来得到它们。use_reentrant=False 是 PyTorch 推荐的新实现。 h = checkpoint(blk, h, use_reentrant=False) if use_ckpt else blk(h) h.sum().backward() flops = {} for use_ckpt in [False, True]: with FlopCounterMode(display=False) as fc: run(use_ckpt) flops[use_ckpt] = fc.get_total_flops() blocks.zero_grad() fwd = 2 * B * dim * dim * L print(f"\n 实测（每层都做检查点）：前向 {fwd:,} FLOPs") print(f" 不做检查点：前向+反向 = {flops[False]:,}（{flops[False] / fwd:.2f} 倍前向）") print(f" 做检查点 ：前向+反向 = {flops[True]:,}（{flops[True] / fwd:.2f} 倍前向） ← 多了一次前向") print(""" 代价：多做一次前向，训练计算量从约 6ND 变成约 8ND（慢约 1/3）。 收益：激活显存大幅下降，可以用更大的批次或更大的模型。 实践中常常只对"激活大、计算便宜"的部分做检查点（选择性重计算），而不是全部重算。""") def demo_budget_table(): """演示 6：不同规模模型在混合精度 + AdamW 下需要多少显存（不含激活）。""" print_header("演示 6：参数 + 梯度 + 优化器状态 的显存（混合精度，16 字节/参数）") for n in [124e6, 1.5e9, 7e9, 70e9, 405e9]: gb = n * 16 / 1e9 print(f" {n / 1e9:>7.3g}B 参数 → {gb:>8,.0f} GB，至少需要 {math.ceil(gb / 80):>3} 块 80 GB 的 GPU") def main(): torch.manual_seed(0) demo_bytes_per_param() demo_adam_state() demo_activation_memory() demo_gradient_accumulation() demo_activation_checkpointing() demo_budget_table() print("\n全部演示运行完毕。") if __name__ == "__main__": main()