S16：训练 Memory Calculator：从参数到账本

内容类型：预习教材（不代表已完成）
日期：2026-12-08
阶段：P2 · AI Systems Engineering 90
总路线：Day 106 / 360
周次 / 节奏：W3 · 周二最小实现
状态：教材已备；学习未完成
主题：memory estimation、optimizer state、activation、peak

一句话定义

训练内存计算器是一份显式假设驱动的容量账本，用于估算不同精度、优化器、并行分片和 batch/context 配置的内存构成，而不是承诺真实峰值。

学习目标

1. 能写一个小函数或表格，分别估算参数、梯度、optimizer 与 activation。
2. 能把 dtype、world size、ZeRO stage、micro-batch 和 context 作为显式输入。
3. 能区分静态估算、框架 allocated/reserved 与设备峰值。
4. 能用差异解释假设缺口，而不是把公式包装成精确容量工具。

核心知识

内存计算必须先声明训练配置。纯 FP32 SGD、BF16 AdamW 和 8-bit optimizer 的每参数状态完全不同；有无 master weights、梯度累积和参数 flattening 也会改变峰值。Calculator 应输出分项与总和，禁止只给一个看似精确的数字。

参数项通常与 num_params × bytes_per_param 成正比；梯度类似；Adam 的一阶、二阶 moment 通常各为 FP32。Activation 更难：它与层结构、需要为 backward 保存的张量、micro-batch、sequence length、checkpointing 及 attention 实现有关。可用一个明确标注为粗估的系数，不应假装适用于所有模型。

分片策略只影响相应状态。ZeRO-1 理想分 optimizer，stage 2 再分 gradient，stage 3 再分 parameter；但每步 all-gather 的临时峰值、通信 bucket 和未分 activation 仍存在。模型加载与 checkpoint 序列化也可能产生短时双份内存。

机制与推导

基础状态：

[ M=P(b_{param}+b_{grad}+b_{master}+b_{m1}+b_{m2})+M_{act}+M_{buffer} ]

理想 ZeRO 分片可按每项除以 shard_factor，例如 stage 2 的 optimizer 和 gradient 除以 world_size，参数保持复制。Activation 粗估可写 M_act ≈ L × B_μ × S × H × c × bytes，其中 c 是与架构/保存张量相关的经验系数，必须标为假设。Gradient accumulation 增加 effective batch：B_eff = B_μ × accumulation × data_parallel_size，它不要求同时保存所有 micro-batch activation。

安全余量可写 capacity_needed = estimated_peak × (1 + headroom)，但 headroom 不能弥补错误模型。

最小练习或观察步骤

1. 用 TS 函数或电子表格定义参数量、dtype、optimizer、layers、micro-batch、sequence、hidden、world size。
2. 输出 parameter/gradient/master/moments/activation/buffer 六行，不只输出 total。
3. 比较 FP32 SGD、BF16 AdamW 与 ZeRO-2 三个纸面配置。
4. 改变 sequence length 为 2 倍，观察 activation 假设怎样变化；注明 attention 可能呈更复杂尺度。
5. 若有 Phase 1 小模型实际峰值记录，仅作对照，写明 estimator 与测量边界。

常见误区与边界

• 用模型文件大小直接推算训练峰值。
• 忽略 master weights 或假设所有 optimizer state 都与参数同 dtype。
• 把 effective batch 当同时驻留设备的 micro-batch。
• 将 ZeRO 理想除法应用到 activation 和所有临时 buffer。
• 报告到 MB 小数点后多位，制造不存在的精确度。

系统场景连接

Capacity 账本能帮助判断一次 LoRA/全参微调是否适合本地或云端、为什么 checkpoint 时 OOM、是否应该先缩短 context 或采用 activation checkpointing。对采购与平台团队，它还能把“需要多少 GPU”拆成可质疑假设，避免只按参数规模估价。本地练习无需真实多 GPU。

自检问题

1. 哪些内存项通常与参数量线性相关，哪些更依赖输入 shape？
2. Gradient accumulation 为什么提高 effective batch 却不同比例增加 activation 峰值？
3. ZeRO-2 理想情况下分片哪些状态？
4. Estimator 与 profiler 数字不一致时应先检查哪些假设？

专业课程对齐

• 阅读 CMU Deep Learning Systems 的 tensor memory、autodiff 与 execution 内容，把 backward 保存值映射到 activation 账本。
• 阅读 PyTorch FSDP 官方文档 的 sharding strategy 与参数状态说明，核对不同阶段哪些状态可能被分片及峰值 gather 边界。
• 阅读 PyTorch Profiler 官方文档 的内存 profiling 能力，区分公式估算、allocator 记录和运行时测量。

深入学习提示

先实现透明而粗糙的分项表，再谈准确率。每一行必须写单位、dtype 和分片因子；activation 系数旁明确“架构相关”。尝试用两个配置得到相同 total 却不同风险：一个状态大、一个 activation 大。前者适合分片，后者可能更适合 checkpointing 或缩 context，这能训练从账本到策略的推理。

学后填写区

• 实际实现方式：
• 输入假设：
• 三个配置的分项结果：
• 与实际测量的差异（如有）：
• 最不可靠的一项估算：