""" ============================================================================= 第 9 讲 · 代码 1：训练时显存都花在哪里——参数、梯度、优化器状态、激活 ============================================================================= 对应讲义：lectures/lecture09_distributed_training.md 第 2 节、第 4.1 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 一个 1750 亿参数的模型，16 位存下来是 350 GB——但训练它需要的显存远不止于此。多出来的是什么？ 演示 1：每个参数要几个字节（混合精度 + Adam：16 字节）；几个真实模型的"模型状态"有多大。 演示 2：在第 1 讲的 base model 上实测：参数、梯度、AdamW 状态各占多少字节——和公式对上。 演示 3：ZeRO 的三个阶段：把优化器状态、梯度、参数切到 N_d 张卡上，每张卡还剩多少（复现 ZeRO 论文图 1）。 演示 4：激活：用 PyTorch 的钩子数出 base model 为反向传播存下了多少字节，看它怎样随 batch、序列长度、 头数、宽度变化；再看激活重算（activation checkpointing）省了多少显存、多花了多少计算。 演示 5：大模型的激活：用 Korthikanti et al. (2022) 的公式算 GPT-3 和 GPT-2 1.5B； 和 ZeRO 论文的粗略估计对照，看注意力矩阵（第 8 讲的 N×N）占了多少。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_memory_accounting.py # 约 3 秒；需要 numpy、PyTorch 依赖本目录的 dist_sim.py 和 code/lecture01 的 tiny_gpt.py、console.py（只建随机初始化的模型，不读缓存）。 【PyTorch 小知识】 - torch.autograd.graph.saved_tensors_hooks(pack, unpack)：在 with 块里，autograd 每次为反向传播"存下" 一个张量时都会调用 pack(张量)。我们在 pack 里记下它的大小——这就是"激活占的显存"。 权重本身也会被存（矩阵乘法的反向要用），我们按 data_ptr()（数据在内存里的地址）把它们排除掉； 同一个张量被存好几次只算一次。 - torch.utils.checkpoint.checkpoint(函数, 输入, use_reentrant=False)：前向时不存这个函数内部的激活， 反向时重新算一遍前向。use_reentrant=False 是 PyTorch 推荐的新实现（不写会有警告）。 - optimizer.state[参数]：AdamW 为每个参数存的状态，包括 exp_avg（一阶矩）和 exp_avg_sq（二阶矩）。 """ import os import sys import numpy as np import torch import torch.nn.functional as F from torch.utils.checkpoint import checkpoint from torch.utils.flop_counter import FlopCounterMode HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from dist_sim import (BYTES_MIXED_ADAM, activation_bytes_per_layer, fmt_bytes, model_state_bytes, # noqa: E402 zero_bytes_per_device) from tiny_gpt import BASE_CONFIG, GPTConfig, count_parameters, make_model # noqa: E402 A100_80 = 80e9 # ============================================================================= # 演示 1：每个参数 16 字节 # ============================================================================= def demo1_bytes_per_param(): print_header("演示 1：混合精度 + Adam，每个参数要存 16 个字节") print(" fp16（或 bf16）参数 2 字节 前向、反向用它") print(" fp16 梯度 2 字节") print(" fp32 主参数 4 字节 更新在这份上做（第 10 讲：小的更新加到 16 位数上会被吞掉）") print(" Adam 一阶矩 m（fp32） 4 字节") print(" Adam 二阶矩 v（fp32） 4 字节") print(" 合计 = 2 + 2 + 12 = 16 字节（ZeRO 论文把优化器部分记作 K = 12）\n") models = [("GPT-2 1.5B", 1.5e9), ("GPT-3 175B", 175e9), ("BLOOM 176B", 176e9), ("PaLM 540B", 540e9)] print(f"{pad('模型', 14)}{pad('16 位参数', 14, '>')}{pad('模型状态 16Ψ', 16, '>')}{pad('至少几张 80 GB 的卡', 22, '>')}") for name, psi in models: st = model_state_bytes(psi) n_cards = f"{st['total'] / A100_80:.1f}" # 先算成字符串，避免 f-string 里嵌套同种引号 print(f"{pad(name, 14)}{pad(fmt_bytes(st['param']), 14, '>')}{pad(fmt_bytes(st['total']), 16, '>')}{pad(n_cards, 22, '>')}") bloom = 176e9 print(f"\nBLOOM 的检查点：bf16 参数 2 + fp32 主参数 4 + Adam 两个矩 8 = 14 字节/参数 → {fmt_bytes(14 * bloom)}" f"（= {14 * bloom / 2 ** 40:.2f} TiB）；只存 bf16 参数 {2 * bloom / 2 ** 30:.0f} GiB") print("（训练日志：完整检查点 2.3TB，只存 bf16 参数 329GB）") print("\n观察 Observations：") print(" 1. ZeRO 论文的例子：GPT-2 1.5B 的 16 位参数只有 3 GB，但训练它至少要 24 GB——8 倍。") print(" 2. GPT-3 的模型状态 2.8 TB，至少要 35 张 80 GB 的卡才放得下——还没算激活。") print(" 单靠'每张卡存一整份模型'的数据并行，根本训练不了这样的模型：必须把模型状态切开（§3–§5）。") print(" 3. 优化器状态（12 字节）是大头，比参数和梯度加起来（4 字节）还多 3 倍——ZeRO 第一步就切它。") # ============================================================================= # 演示 2：在 base model 上实测 # ============================================================================= def demo2_measure_states(): print_header("演示 2：第 1 讲的 base model（fp32 训练）实测：参数、梯度、AdamW 状态") model = make_model(BASE_CONFIG, seed=0) opt = torch.optim.AdamW(model.parameters(), lr=1e-3) idx = torch.randint(0, 44, (4, 96), generator=torch.Generator().manual_seed(0)) loss = F.cross_entropy(model(idx).view(-1, 44), idx.view(-1)) loss.backward() opt.step() n = count_parameters(model) p_bytes = sum(p.numel() * p.element_size() for p in model.parameters()) g_bytes = sum(p.grad.numel() * p.grad.element_size() for p in model.parameters()) s_bytes = sum(v.numel() * v.element_size() for st in opt.state.values() for v in st.values() if torch.is_tensor(v) and v.dim() > 0) print(f"参数个数 Ψ = {n:,}") print(f" 参数 {p_bytes:>12,} 字节 = {p_bytes / n:.0f} 字节/参数") print(f" 梯度 {g_bytes:>12,} 字节 = {g_bytes / n:.0f} 字节/参数") print(f" AdamW 状态 {s_bytes:>12,} 字节 = {s_bytes / n:.0f} 字节/参数（exp_avg + exp_avg_sq）") print(f" 合计 {p_bytes + g_bytes + s_bytes:>12,} 字节 = {(p_bytes + g_bytes + s_bytes) / n:.0f} 字节/参数") print("\n观察 Observations：") print(" 1. 纯 fp32 训练也是每个参数 16 字节（4 + 4 + 8），只是分法不同：没有单独的'主参数'，参数本身就是 fp32。") print(" 混合精度并不省模型状态的显存（还是 16 字节），省的是激活和计算时间（第 8 讲 §3.4，第 10 讲）。") # ============================================================================= # 演示 3：ZeRO # ============================================================================= def demo3_zero(): print_header("演示 3：ZeRO——把模型状态切到 N_d 张数据并行的卡上（每张卡的字节数）") psi = 7.5e9 names = ["普通数据并行", "ZeRO-1 切优化器状态", "ZeRO-2 再切梯度", "ZeRO-3 再切参数"] print(f"Ψ = 7.5B（ZeRO 论文图 1 的例子），K = 12") print(f"{pad('', 22)}" + "".join(pad(f"N_d = {n}", 14, ">") for n in (8, 64, 1024))) for stage, name in enumerate(names): print(f"{pad(name, 22)}" + "".join(pad(fmt_bytes(zero_bytes_per_device(psi, n, stage)), 14, ">") for n in (8, 64, 1024))) print("\n80 GB 的卡、只看模型状态（不算激活），最多能训练多大的模型：") print(f"{pad('', 22)}" + "".join(pad(f"N_d = {n}", 14, ">") for n in (8, 64, 1024))) for stage, name in enumerate(names): cells = [] for n in (8, 64, 1024): per_param = zero_bytes_per_device(1.0, n, stage) cells.append(pad(f"{A100_80 / per_param / 1e9:.1f}B", 14, ">")) print(f"{pad(name, 22)}" + "".join(cells)) print("\n观察 Observations：") print(" 1. N_d = 64 时复现了 ZeRO 论文图 1：120 GB → 31.4 GB → 16.6 GB → 1.88 GB。") print(" 2. ZeRO-1、ZeRO-2 有天花板：参数（2Ψ）或参数 + 梯度（4Ψ）仍然每张卡一整份，卡再多也省不掉；") print(" ZeRO-1 在 N_d = 1024 时也只能到 20B 左右。ZeRO-3 的显存随卡数线性下降——1024 张卡能放下 5 万亿参数的模型状态。") print(" 3. 天下没有免费的午餐：ZeRO-3 每一步都要先把参数 all-gather 回来才能算，通信量是普通数据并行的 1.5 倍（代码 02）。") # ============================================================================= # 演示 4：激活 # ============================================================================= def saved_activation_bytes(model, B, T, use_checkpoint=False): """ 数一数一次前向（+ 交叉熵）为反向传播存下了多少字节的激活（不含参数本身）。 【思路】用 saved_tensors_hooks 拦截 autograd 存下的每一个张量，按 (地址, 形状, 类型) 去重， 排除模型参数，把剩下的字节数加起来。 【参数】model：TinyGPT；B, T：batch 大小、序列长度；use_checkpoint：True 时每个 Block 用激活重算 【返回】字节数（整数） 【边界】只算"为反向存下的"张量；前向中途临时分配、马上释放的不算。所以这是激活显存的下界。 """ params = {p.data_ptr() for p in model.parameters()} seen = {} def pack(t): if t.data_ptr() not in params: seen[(t.data_ptr(), tuple(t.shape), t.dtype)] = t.numel() * t.element_size() return t idx = torch.randint(0, model.cfg.vocab_size, (B, T), generator=torch.Generator().manual_seed(0)) with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): logits = forward_maybe_checkpointed(model, idx, use_checkpoint) F.cross_entropy(logits.reshape(-1, logits.size(-1)), idx.reshape(-1)) return sum(seen.values()) def forward_maybe_checkpointed(model, idx, use_checkpoint): """TinyGPT 的前向；use_checkpoint=True 时每个 Block 包一层 checkpoint（只存 Block 的输入）。""" T = idx.shape[1] x = model.tok_emb(idx) + model.pos_emb(torch.arange(T)) for blk in model.blocks: x = checkpoint(blk, x, use_reentrant=False) if use_checkpoint else blk(x) return model.head(model.ln_f(x)) def train_step_flops(model, B, T, use_checkpoint): """前向 + 反向一共多少 FLOPs（PyTorch 的 FLOP 计数器）。""" idx = torch.randint(0, model.cfg.vocab_size, (B, T), generator=torch.Generator().manual_seed(0)) with FlopCounterMode(display=False) as fc: logits = forward_maybe_checkpointed(model, idx, use_checkpoint) F.cross_entropy(logits.reshape(-1, logits.size(-1)), idx.reshape(-1)).backward() model.zero_grad() return fc.get_total_flops() def demo4_activations(): print_header("演示 4：base model 为反向传播存下的激活（fp32，2 层），以及激活重算") rows, X, y = [], [], [] configs = [(64, 4, 8, 96), (64, 4, 16, 96), (64, 4, 8, 48), (64, 4, 8, 24), (128, 4, 8, 96), (64, 8, 8, 96), (128, 8, 16, 48)] print(f"{pad('d', 6, '>')}{pad('头数 a', 8, '>')}{pad('batch b', 9, '>')}{pad('序列 s', 8, '>')}" f"{pad('存下的激活', 16, '>')}{pad('用激活重算', 14, '>')}{pad('省了', 8, '>')}") for d, a, b, s in configs: cfg = GPTConfig(vocab_size=44, block_size=96, n_layer=2, n_head=a, d_model=d) model = make_model(cfg, seed=0) full = saved_activation_bytes(model, b, s) ck = saved_activation_bytes(model, b, s, use_checkpoint=True) print(f"{pad(d, 6, '>')}{pad(a, 8, '>')}{pad(b, 9, '>')}{pad(s, 8, '>')}{pad(f'{full:,}', 16, '>')}" f"{pad(f'{ck:,}', 14, '>')}{pad(f'{full / ck:.1f}×', 8, '>')}") L = cfg.n_layer X.append([L * s * b * d, L * a * s * s * b, s * b]) # 三个特征：每层每 token 每维、每层每头每对位置、每 token y.append(full / 4) # fp32：字节 ÷ 4 = 存下的数的个数 coef, *_ = np.linalg.lstsq(np.array(X, dtype=float), np.array(y), rcond=None) pred = np.array(X, dtype=float) @ coef print(f"\n最小二乘拟合（单位：存下的数的个数）：激活 ≈ {coef[0]:.2f}·L·s·b·h + {coef[1]:.2f}·L·a·s²·b + {coef[2]:.1f}·s·b" f" （7 个配置的最大相对误差 {np.max(np.abs(pred - y) / y):.1e}）") print("Korthikanti et al. 的 GPT 层（16 位，字节）：34·s·b·h + 5·a·s²·b → 折合约 17 个数 / (层·token·维) + 2.5 个 / (层·头·位置对)") model = make_model(BASE_CONFIG, seed=0) f0, f1 = train_step_flops(model, 32, 96, False), train_step_flops(model, 32, 96, True) print(f"\nbase model（batch 32）前向 + 反向的 FLOPs：不重算 {f0:,}，每个 Block 重算 {f1:,}，多 {f1 / f0 - 1:.1%}") print("\n观察 Observations：") print(" 1. 激活随 batch 线性增长（batch 8 → 16 正好翻倍），随序列长度**超线性**增长：") print(" 三个特征的拟合几乎完美——每层每个 token、每个隐藏维度存 17 个数，每层每个头每一对位置存 2 个数") print(" （注意力的 s×s 矩阵：softmax 的输出等）。头数从 4 加到 8，只有 s² 那一项变大。") print(" 2. 第一项的系数 17.00 和 Korthikanti 的 34 字节（16 位下 17 个数）正好一样——同一种结构，同样的账。") print(" 3. 激活重算只存每个 Block 的输入，激活省了 9–17 倍（base model 的配置是 12.3 倍）；代价是反向前把 Block 的前向再算一遍，") print(" 这个模型多了 23.9% 的 FLOPs（理论上限是 1/3：前向 2N + 反向 4N + 重算 2N = 8N vs 6N；这里嵌入、输出层和交叉熵不重算）。") print(" 这就是第 8 讲 §7 HFU > MFU 的来源。") # ============================================================================= # 演示 5：大模型的激活 # ============================================================================= def demo5_big_activations(): print_header("演示 5：大模型的激活（Korthikanti et al. 2022 的公式，16 位）") s, b, h, a, L = 2048, 1, 12288, 96, 96 print(f"GPT-3 175B：{L} 层，h = {h}，a = {a}，一个序列 s = {s}（micro-batch b = {b}）") rows = [("不并行、不重算", activation_bytes_per_layer(s, b, h, a)), ("张量并行 t = 8", activation_bytes_per_layer(s, b, h, a, t=8)), ("张量并行 t = 8 + 序列并行", activation_bytes_per_layer(s, b, h, a, t=8, sequence_parallel=True)), ("完全重算（只存每层输入）", activation_bytes_per_layer(s, b, h, a, recompute=True))] print(f"{pad('做法', 30)}{pad('每层', 12, '>')}{pad(f'{L} 层合计', 14, '>')}") for name, v in rows: print(f"{pad(name, 30)}{pad(fmt_bytes(v), 12, '>')}{pad(fmt_bytes(v * L), 14, '>')}") att_share = 5 * a * s / h / (34 + 5 * a * s / h) print(f"其中注意力的 s² 项（5as/h = {5 * a * s / h:.0f}）占每层激活的 {att_share:.0%}") s, b, h, a, L = 1024, 32, 1600, 25, 48 full = activation_bytes_per_layer(s, b, h, a) * L no_att = 34 * s * b * h * L zero_est = 12 * s * b * h * L * 2 recomp = activation_bytes_per_layer(s, b, h, a, recompute=True) * L print(f"\nGPT-2 1.5B（{L} 层，h = {h}，a = {a}），序列 1K、batch 32——ZeRO 论文的例子：") print(f" ZeRO 论文的估计（12·h·s·b·层数 个 16 位数） {fmt_bytes(zero_est)} （论文：约 60 GB）") print(f" Korthikanti 公式，含注意力的 s² 项 {fmt_bytes(full)}") print(f" Korthikanti 公式，去掉 s² 项（相当于用 FlashAttention） {fmt_bytes(no_att)}") print(f" 完全重算（每层只存 2sbh 字节） {fmt_bytes(recomp)} （ZeRO 论文：重算后约 8 GB）") print("\n观察 Observations：") print(" 1. GPT-3 的一个 2048 token 的序列，不并行、不重算，每层的激活就要 2.87 GB、96 层 275 GB——") print(" 比 GPT-3 的 16 位参数（350 GB）只少一点。其中 70% 是注意力的 s×s 矩阵。") print(" 2. 张量并行把大部分激活也切成 t 份；加上序列并行（把 LayerNorm、dropout 那些部分沿序列切开）之后整体除以 t。") print(" 完全重算只存每层的输入：每层 50 MB，但要多算一遍前向。") print(" 3. 同一个 GPT-2 1.5B 的例子，三种记账相差近 5 倍：ZeRO 论文的 60 GB 是粗略估计（每层每 token 每维 12 个数），") print(" 没有算注意力的 s² 项；把它算上是 287 GB，去掉它（FlashAttention 不存 s×s 矩阵）是 86 GB。") print(" 结论不变——激活和模型状态一样是显存大户——但'激活有多大'很依赖实现，引用数字时要看它怎么算的。") if __name__ == "__main__": demo1_bytes_per_param() demo2_measure_states() demo3_zero() demo4_activations() demo5_big_activations()