Training memory 由什么构成？怎么估？
先弄清楚：每个概念是什么
- Parameter 与 gradient
- weight 和 gradient 各一份，mixed precision 下各 2 bytes。例：14B 模型的 bf16 weight 是 28 GB。
- Optimizer state
- Adam 的 m、v 加一份 fp32 master weight，共 12 bytes，是大头（推导见 optimizer）。
- Activation
- forward 时留给 backward 用的中间结果。和参数量无关，随 \(b\)、\(s\) 线性涨，attention 的 score 随 \(s^2\) 涨。
- Buffer 与 fragmentation
- gradient all-reduce 前常拼成一个大 buffer，1.5B 模型的 fp32 buffer 就 6 GB。fragmentation 严重时，ZeRO 观察到还剩 30% 以上空闲却 OOM。
- LoRA / QLoRA
- 冻结原 weight，只训练低秩的 \(BA\)。QLoRA 再把冻结的 weight 压成 4-bit。
看图：16 bytes 怎么分，activation 有多大
30 秒回答
Four parts. Parameters and gradients in bf16 take 2 bytes each. Adam's state, an fp32 master copy plus two moments, takes 12, so model states total 16 bytes per parameter: a 14B model needs 28 GB for weights but 224 GB to train. Activations saved for backward take about s·b·h·(34 + 5as/h) bytes per layer; TP plus SP divides that by t, and recomputation cuts it further. Then add buffers and fragmentation. LoRA trains only r·(d_in + d_out) parameters per matrix, and QLoRA stores frozen weights in 4-bit.
怎么估：四部分相加
| 部分 | bytes | 14B 例 |
|---|---|---|
| parameter（bf16） | \(2\Psi\) | 28 GB |
| gradient（bf16） | \(2\Psi\) | 28 GB |
| optimizer state（fp32 master、m、v） | \(12\Psi\) | 168 GB |
| model states 合计 | \(16\Psi\) | 224 GB |
| activation | 每层见下式，\(\times L\) | 看 \(b\)、\(s\) |
Activation（fp16，不做 parallelism / recompute）
\[ \text{每层} = sbh\Big(34 + \frac{5as}{h}\Big)\ \text{bytes} \]34 = attention 11 + MLP 19 + 两个 LayerNorm 4（单位 \(sbh\)）；\(5as/h\) 是 \(s\times s\) 的 softmax 输出、dropout mask 和 dropout 输出。
\[ \text{TP + SP：} \div t;\qquad \text{再加 selective recompute：} \frac{34\,sbh}{t} \] \[ \text{full recompute：} 2sbh\ \text{（只存每层输入）} \]PP 用 1F1B 时第一个 stage 存 \(p\) 个 micro-batch × \(L/p\) 层，仍是 \(L\) 层。
手算例子
\(s = 4096\)、\(h = 8192\)、\(a = 64\)、\(L = 80\)、\(b = 1\)：
- \(5as/h = 160\)，每层 \(sbh \times 194 \approx 6.5\) GB，80 层 521 GB，放不下。
- TP = 8 + SP + selective：每层 \(34sbh/8 \approx 143\) MB，80 层 11.4 GB。
- full recompute：每层 \(2sbh \approx 67\) MB，80 层 5.4 GB。
LoRA：每个矩阵加 \(r(d_\text{in} + d_\text{out})\) 个 trainable 参数，\(A\) 随机、\(B = 0\) → 初始 \(\Delta W = 0\)。冻结的 weight 只占 2 bytes/param，16 bytes 只算在 LoRA 参数上。例：\(h = 4096\)、32 层、只加在 \(W_q\)、\(W_v\)、\(r = 8\) → \(2\times32\times8\times8192 \approx 4.2\)M。QLoRA：冻结 weight 存成 NF4，65B finetune 从 >780 GB 降到 <48 GB，细节见 LoRA。
必答点
- 四部分：parameter、gradient、optimizer state、activation（+ buffer、fragmentation）
- mixed precision Adam：2 + 2 + 12 = 16 bytes/param；14B → 224 GB
- activation 每层 \(sbh(34 + 5as/h)\)；TP + SP 除以 \(t\)，selective 去掉 \(5as/h\)，full recompute 只剩 \(2sbh\)
- LoRA \(r(d_\text{in} + d_\text{out})\)，\(B = 0\)；QLoRA = NF4 + double quantization + paged optimizer
追问
SGD、momentum、Adam 各要多少？"4 倍参数量"怎么来？
全 fp32 时 weight + gradient 8 bytes，加上 state 是 SGD 8、momentum 12、Adam 16 bytes/param。Adam 的 w、g、m、v 四份就是"4 倍"，表格见 optimizer。
ZeRO 能省多少？
ZeRO-1/2/3 依次切 optimizer state、gradient、parameter。7.5B、DP \(N = 64\) 时每卡从 120 GB 降到 31.4、16.6、1.9 GB，公式见 DP、ZeRO、FSDP。
LoRA 为什么只把 B 初始化成 0？
\(B = 0\) 让训练从原模型出发。\(A\)、\(B\) 都是 0 时两边的 gradient 全是 0，学不动，式子见 LoRA。
OOM 了怎么办？
先减 micro-batch 加 gradient accumulation，再上 selective recompute，列表见 activation recompute。
来源Korthikanti et al. 2022 §4.1 eq. 1（p.4）、§4.2 eq. 2–5（pp.5–8）、§5 Table 2（p.9）· Rajbhandari et al. 2020, ZeRO Fig. 1（p.3）、§3.1–3.2（pp.7–8）· Hu et al. 2021, LoRA §4.1–4.2（pp.4–5）· Dettmers et al. 2023, QLoRA §1（p.2）、§3（pp.4–5）