v1 单卡上的三招：BF16、activation checkpointing、FlashAttention
v0 的 model state 128 GB，一条 8K 序列的 activation 约 430 GB，fp32 还算得慢。v1 先不加卡，只用三招。BF16 让 matmul 跑上 Tensor Core，FlashAttention 和 activation checkpointing 把 activation 压到 2.4 GB，但 128 GB 的 model state 一点没少。
第一招 BF16：先看一个浮点数怎么存
浮点数就是二进制的科学计数法，比如 \(1.5 \times 2^{3}\)。它分三段：
- sign：1 bit，管正负
- exponent：2 的几次方。bit 越多，能表示的数越大也越小，这叫 range，像尺子有多长
- mantissa：小数点后那部分。bit 越多，相邻两个数挨得越近，这叫 precision，像刻度有多密
| 格式 | sign / exponent / mantissa | 最大值 | 1 附近相邻两数的间距 | 一句话 |
|---|---|---|---|---|
| FP32 | 1 / 8 / 23 | 3.4e38 | \(2^{-23}\) ≈ 1.2e-7 | 尺子长、刻度密，但占 4 bytes |
| FP16 | 1 / 5 / 10 | 65504 | \(2^{-10}\) ≈ 0.001 | 刻度够密，尺子太短 |
| BF16 | 1 / 8 / 7 | 3.4e38 | \(2^{-7}\) ≈ 0.008 | 尺子长，刻度粗 |
为什么大模型都选 BF16：
- 只占 2 bytes，能跑满 Tensor Core：H100 上 BF16 是 989 TFLOPS，纯 FP32 只有 67 TFLOPS，差约 15 倍
- range 和 FP32 一样，所以不需要 loss scaling。代价是刻度粗，靠下面的 fp32 master weight 和 fp32 accumulation 补回来
FP16 为什么要 loss scaling
gradient 常常很小，比如 1e-8，FP16 存不下，直接变成 0（underflow）。所以 FP16 要在 backward 之前把 loss 乘一个大数 S（比如 \(2^{16}\)），让所有 gradient 一起放大，optimizer 更新之前再转成 fp32 除以 S。S 太大又会超过 65504 变成 inf，所以实际用 dynamic 的。BF16 把这一套都省了。
Mixed precision 的配方
“mixed”的意思是，怕误差的地方用 fp32，其余都用 bf16。
- fp32 master weight 是“正本”，每一步复制出一份 bf16 副本
- forward、backward 的 matmul 输入是 bf16，Tensor Core 内部用 fp32 accumulation（累加）。softmax、LayerNorm 这类要求和的操作也用 fp32
- 得到 bf16 的 gradient，转成 fp32，更新 fp32 master weight
两处 fp32 防的是同一件事：大数加小数，刻度太粗，小数就被吃掉。
- master weight：w = 1、update 是 1e-4 时，BF16 在 1 附近的刻度是 0.008，1 + 0.0001 舍入后还是 1，update 丢了。fp32 的刻度是 1e-7，留得住
- accumulation：h = 4096 的 dot product 要加 4096 个乘积，和越加越大，后面的小项在 bf16 里同样被吃掉
mixed precision 不减 model state，每个 parameter 还是 2 + 2 + 4 + 4 + 4 = 16 bytes。它省的是 activation（每个数 4 → 2 bytes，约 860 → 430 GB）和时间（matmul 跑上 Tensor Core，memory-bound 的算子要搬的 bytes 减半）。
第二招 FlashAttention：不把 s × s 写回 HBM
430 GB 里 90% 是 attention 的 s × s。标准 attention 每个 head 先算 score \(S = QK^\top\)，形状是 s × s：
- s = 8192 时是 8192² ≈ 6700 万个数，bf16 下每个 head 134 MB。32 个 head 加上 softmax 的输出和 dropout，每层约 10.7 GB
- 它还慢：S 写回 HBM，读出来做 softmax，结果再写回，再读出来乘 V，时间全花在搬运上（第 0 课的 memory-bound）
FlashAttention 的做法：
- tiling：把 Q、K、V 切成小块搬到案板（SRAM）上，每次只算一小块 score，乘上 V 累加进输出就扔掉
- online softmax：softmax 要整行的最大值和总和，可一次只看得到一块。于是边看边记当前的最大值 m 和总和，m 变大时把之前的结果乘上 \(e^{m_\text{old} - m_\text{new}}\) 校正。结果和一次算整行完全一样，不是近似
- backward：不存 score，用 Q、K、V 分块重算
- 省多少：去掉 \(5as/h\) 那项，每层 11.9 → 1.14 GB，36 层 430 → 41 GB
第三招 activation checkpointing：只存每层的输入
- 思路：forward 时每层只留输入（checkpoint），其余中间结果用完就扔。backward 走到这一层时，从 checkpoint 重跑这一层的 forward，拿回中间结果
- 省多少：每层输入是 [b, s, h]，2sbh = 2 × 8192 × 4096 bytes ≈ 67 MB，36 层 ≈ 2.4 GB，也就是 41 → 2.4 GB。backward 时临时重建一层的 1.14 GB，峰值约 3.6 GB
- 代价：多跑一次 forward，FLOPs 从 2 + 4 = 6 变成 2 + 4 + 2 = 8，多约 1/3，MFU 也跟着降
- Megatron 的 selective recompute：只重算 attention 里 s × s 的那几步（\(5as/h\) 那项），GPT-3 上省 70% 的 activation，只多约 2.7% 的 FLOPs。FlashAttention 本来就不存这部分，开了以后要再往下压只能靠 full recompute
OOM 工具箱：按代价从小到大
| 办法 | 省什么 | 代价 |
|---|---|---|
| BF16 mixed precision | activation 减半，还更快 | 几乎没有 |
| FlashAttention | attention 的 s × s：430 → 41 GB | 没有，还更快 |
| 小 micro-batch + gradient accumulation | activation 随 b 线性下降 | matmul 变小，GPU 吃不饱 |
| activation checkpointing | 41 → 2.4 GB | 多约 1/3 compute |
| ZeRO / FSDP、TP、PP、CP（第 4–7 课） | 把 model state 和 activation 切到多张卡 | communication |
| CPU offload | 把 optimizer state 放到 CPU memory | 走 PCIe，很慢 |
去查一下你在 TikTok 用的 DeepSpeed + HF 里，这三招各是一个开关：DeepSpeed config 里的 "bf16": {"enabled": true}；HF 的 gradient_checkpointing=True（或 model.gradient_checkpointing_enable()）；from_pretrained(..., attn_implementation="flash_attention_2")。翻一下当时的 config 和启动脚本，看开了哪几个，sequence length 设的多少。如果 attention 写的是 "sdpa"，PyTorch 也可能在内部选用 FlashAttention 的 kernel。面试官很可能追问：“Why did you turn on gradient checkpointing, and what did it cost?”
30 秒英文回答
Before adding GPUs, use three tricks. BF16 mixed precision runs matmuls in BF16 on Tensor Cores, with FP32 accumulation and FP32 master weights. BF16 keeps FP32's eight exponent bits, so unlike FP16 it needs no loss scaling. It halves activations and speeds things up, but model state stays at 16 bytes per parameter. FlashAttention computes attention tile by tile in SRAM with an online softmax, so the s-by-s scores never reach HBM: for one 8K sequence on an 8B model, activations drop from about 430 to 41 GB. Activation checkpointing keeps only each layer's input and recomputes the rest in backward: 2.4 GB, for about a third more compute.
自测
为什么大模型训练用 BF16，不用 FP16？
BF16 有 8 bit exponent，range 和 FP32 一样，不用 loss scaling。代价是 mantissa 只有 7 bit，靠 fp32 master weight 和 fp32 accumulation 补回来。
开了 BF16 mixed precision，8B 模型的 model state 是多少？
还是 128 GB，因为 bf16 省下的 4 bytes 被 fp32 master weight 吃回去。省的是 activation 和时间。
FlashAttention 没有减少 FLOPs，为什么还更快？
attention 是 memory-bound，FlashAttention 省掉的是 s × s 在 HBM 的读写。论文里（A100、GPT-2 medium）HBM 读写从 40.3 GB 降到 4.4 GB，forward + backward 从 41.7 ms 降到 7.3 ms。
full recompute 和 selective recompute 有什么区别？开了 FlashAttention 还要 selective 吗？
full 整层重跑 forward，多约 1/3 compute。selective 只重算 attention 里 s × s 的那几步，FLOPs 只多几个百分点。开了 FlashAttention 就不用 selective，要继续压只能用 full（41 → 2.4 GB）。
显存不够有哪些办法？按代价排。
先用不掉速的（BF16、FlashAttention），再花 compute（小 micro-batch 加 gradient accumulation、activation checkpointing），最后花通信（ZeRO/FSDP、TP/PP/CP、CPU offload）。
想看细一点 6.15 mixed precision（dynamic loss scaling、FP8）· 6.14 activation recompute（\(\sqrt n\) checkpointing、selective 的数）· 6.25 FlashAttention（online softmax 公式、shape、NumPy 代码）· 6.20 gradient accumulation