一句话：在不改动模型质量的前提下，用「精度换字节、算力换显存、并行换等待」三类手段，把更大的模型/批量塞进有限的显存并榨干算力。 关键年份：Mixed Precision Training（Micikevicius et al. 2017，arXiv:1710.03740）；Sublinear Memory / 梯度检查点（Chen et al. 2016，arXiv:1604.06174）；FlashAttention（Dao et al. 2022，arXiv:2205.14135）；8-bit Optimizers（Dettmers et al. 2021，arXiv:2110.02861）。 前置阅读：Transformer 架构、训练系统总览、推理优化
训练大模型时，单卡显存几乎永远是第一约束。要理解优化从哪里下手，先看显存被谁吃掉。以 Adam + 混合精度训练一个参数量为
- 模型权重：fp16 副本
，外加 fp32 主权重 ； - 优化器状态：Adam 的一阶/二阶动量各
（fp32），共 ； - 梯度：
（fp16）或 ； - 激活值（activations）：与 batch size、序列长度、层数线性相关，长序列下往往是最大头。
前三项合计约
混合精度训练（Mixed Precision）
把前向/反向的计算放到 16 位浮点（fp16 或 bf16），同时保留一份 fp32 主权重用于参数更新，这就是混合精度训练的核心（Micikevicius et al. 2017）。收益有两点：16 位张量减半激活与梯度显存；现代 GPU 的 Tensor Core 对 16 位矩阵乘有数倍吞吐优势。
fp16 与 bf16 的关键差别在动态范围：
| 格式 | 位宽 | 指数位 | 尾数位 | 动态范围 | 是否需 loss scaling |
|---|---|---|---|---|---|
| fp32 | 32 | 8 | 23 | 很宽 | — |
| fp16 | 16 | 5 | 10 | 窄，小梯度易下溢 | 通常需要 |
| bf16 | 16 | 8 | 7 | 与 fp32 同宽 | 一般不需要 |
fp16 只有 5 位指数，反向传播中大量小梯度会下溢为 0。论文给出的解法是 loss scaling：把损失乘以一个放大因子
使小梯度移入 fp16 可表示区间；更新前再除回
PyTorch 的 torch.cuda.amp（Automatic Mixed Precision，AMP）把上述逻辑封装为 autocast + GradScaler：autocast 自动决定哪些算子走低精度（矩阵乘、卷积）、哪些保留 fp32（如 softmax、layernorm 的归约），GradScaler 负责动态 loss scaling。
激活重计算 / 梯度检查点（Gradient Checkpointing）
激活值在前向时被缓存，供反向计算梯度使用，显存随层数线性增长。梯度检查点（Chen et al. 2016）的思路是只保存少量检查点激活，其余在反向时按需重算，用一次额外前向换取显存下降。
对
Selective recomputation（选择性重计算） 是更精细的版本：只重算那些「计算便宜但占显存大」的算子（如注意力中间结果），保留「计算昂贵」的部分。Megatron-LM 的 selective activation recomputation 即按此思路把重算开销压到个位数百分比，是大模型训练的常用配置。
梯度累积（Gradient Accumulation）
当目标 batch size 装不进单卡时，把它拆成
for i, batch in enumerate(loader):
loss = model(batch) / accum_steps # 缩放，保证梯度等价于大 batch 平均
loss.backward() # 梯度自然累加
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()它不省静态显存（权重/优化器状态照旧），但避开了大 batch 激活同时驻留的峰值，让小卡也能复现大 batch 的训练超参。代价是更新频率下降、吞吐略降。
FlashAttention：省激活显存的注意力
标准注意力会显式构造
8-bit 优化器
Adam 的两个动量状态合计 bitsandbytes 提供，号称两行代码即可替换原优化器。
通信与计算重叠
进入多卡/集群后，瓶颈从「显存」转向「等待通信」。核心优化是让通信藏到计算背后：
- 梯度反向与 all-reduce 重叠：反向是从后往前逐层算梯度，靠后的层先算完，可立即触发该层（或一桶 bucket）的梯度 all-reduce，与前面层的反向计算并行；DDP 的 gradient bucketing 即此机制。
- ZeRO / FSDP 的参数通信重叠：分片参数在用到前异步 all-gather（prefetch），用完即释放，把权重收集藏到上一层计算里。
- 流水线并行的 micro-batch 重叠：用 1F1B 等调度让不同 micro-batch 的前后向交错，压缩 pipeline bubble。
这些机制的目标一致：让有效算力利用率（MFU）尽量接近通信不可见的理想值。具体并行策略（DP/TP/PP/ZeRO）见训练系统总览。
小结：按瓶颈选工具
| 瓶颈 | 首选手段 | 性质 |
|---|---|---|
| 激活显存（长序列/大 batch） | FlashAttention、梯度检查点、selective recompute | 算力换显存 |
| 静态显存（权重/优化器） | 混合精度、8-bit 优化器、ZeRO 分片 | 字节换精度/通信 |
| 凑不出大 batch | 梯度累积 | 时间换等效 batch |
| 多卡等待 | 通信-计算重叠 | 并行换等待 |
实践上这些手段叠加使用：bf16 + FlashAttention + 选择性重计算 + ZeRO 几乎是现代 LLM 预训练的标配组合。优化前先用显存分析定位是激活还是静态开销占主导，再对症下药。
参考文献
- Micikevicius et al. Mixed Precision Training. ICLR 2018. arXiv:1710.03740
- Chen et al. Training Deep Nets with Sublinear Memory Cost. 2016. arXiv:1604.06174
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
- Dettmers et al. 8-bit Optimizers via Block-wise Quantization. ICLR 2022. arXiv:2110.02861