---

一、排查思路

训练大模型时出现 OOM（Out Of Memory，显存不足），是最常见也最让人头疼的问题之一。

排查和优化的核心思路是：先定位显存花在哪里，再按代价从低到高逐个优化。

显存主要花在四部分：

• 模型参数
• 梯度
• 优化器状态
• 中间激活值

不同部分的优化手段不同，代价也不同。优先级应该是：

• 1、先确认是不是真的显存不足，而不是其他问题；
• 2、从代价最低的手段开始，比如减小 batch size、开启混合精度；
• 3、再考虑梯度累积、梯度检查点；
• 4、然后是 ZeRO、并行策略；
• 5、最后才是模型结构层面的改动。

下面按优先级展开。

---

二、第一步：先定位问题

1、确认 OOM 的具体位置

OOM 可能发生在：

• 模型加载时；
• 前向传播时；
• 反向传播时；
• 优化器更新时；
• 保存 checkpoint 时。

不同位置说明瓶颈不同：

• 加载时 OOM：模型参数太大；
• 前向 OOM：激活值太大；
• 反向 OOM：激活值 + 梯度；
• 优化器更新 OOM：优化器状态太大；
• 保存时 OOM：显存汇聚。

2、看显存组成

用工具观察显存：

• nvidia-smi：看整体占用；
• torch.cuda.memory_summary()：看 PyTorch 分配详情；
• PyTorch Profiler：看各操作显存；
• Nsight Systems：看时间线和显存峰值。

重点区分：

• 模型状态显存：参数、梯度、优化器状态；
• 激活值显存：中间结果；
• 临时显存：计算过程中的临时 buffer；
• 碎片：显存碎片导致的浪费。

3、估算理论显存

假设模型有 $\Psi$ 个参数，混合精度 + Adam：

• FP16 参数：$2\Psi$
• FP16 梯度：$2\Psi$
• FP32 主权重：$4\Psi$
• FP32 一阶动量：$4\Psi$
• FP32 二阶动量：$4\Psi$

模型状态合计约：

$$ 16\Psi $$

激活值约为：

$$ \text{Activation} \approx L \times b \times s \times d_{\text{model}} \times \text{bytes} $$

其中 $L$ 是层数，$b$ 是 batch size，$s$ 是序列长度，$d_{\text{model}}$ 是隐藏维度。

如果理论显存已经超过单卡容量，说明必须用显存优化或并行策略。

---

三、第二步：按优先级优化

优先级 1：减小 batch size

代价最低，最先尝试。

• 减小 micro-batch size；
• 如果 batch size 已经为 1，考虑梯度累积；
• 注意 batch size 太小可能影响训练稳定性。

优点：立刻见效，不需要改代码结构。 缺点：吞吐下降，训练变慢；batch 太小可能影响收敛。

优先级 2：开启混合精度

代价低，收益大。

• 用 FP16 或 BF16 代替 FP32；
• 参数和激活值显存减半；
• BF16 动态范围大，更稳定。

优点：显存减半，速度提升。 缺点：FP16 需要 Loss Scaling；某些操作需要保持 FP32。

优先级 3：梯度累积

用时间换显存。

• 把大 batch 拆成多个 micro-batch；
• 每个 micro-batch 前向反向，累积梯度；
• 累积若干步后再更新参数。

这样等效 batch size 不变，但显存只需要一个小 micro-batch。

优点：显存大幅下降，等效 batch 不变。 缺点：训练变慢；BatchNorm 等对 batch 敏感的操作可能受影响。

优先级 4：梯度检查点（Activation Checkpointing）

用计算换显存。

• 前向不保存全部激活值；
• 反向时从检查点重算；
• 显存从 $O(L)$ 降到 $O(\sqrt{L})$。

优点：激活值显存大幅下降。 缺点：计算量增加约 20%–30%。

适合：显存瓶颈在激活值，计算资源相对充足。

优先级 5：优化器显存优化

• 使用 ZeRO-1 切分优化器状态；
• 使用 8-bit Adam 等低精度优化器；
• 使用 Adafactor 等省显存优化器；
• 使用 CPU Offload，把优化器状态放 CPU。

优点：优化器状态显存大幅下降。 缺点：可能增加通信或降低速度。

优先级 6：ZeRO 系列

• ZeRO-1：切分优化器状态；
• ZeRO-2：再切分梯度；
• ZeRO-3：再切分参数。

每卡显存：

• ZeRO-1：约 $4\Psi + 12\Psi/N$
• ZeRO-2：约 $2\Psi + 14\Psi/N$
• ZeRO-3：约 $16\Psi/N$

优点：显存大幅下降，支持更大模型。 缺点：通信开销增加，尤其 ZeRO-3。

优先级 7：并行策略

• 张量并行 TP：切分单层参数，适合节点内；
• 流水线并行 PP：切分层，适合跨节点；
• 数据并行 DP：复制模型，提高吞吐；
• 3D 并行：组合使用。

优点：能支持超大模型。 缺点：实现复杂，通信开销大。

优先级 8：模型结构优化

• 使用 GQA/MQA 减少 KV 头；
• 使用 MoE 稀疏激活；
• 减少层数或隐藏维度；
• 使用更高效的注意力。

优点：从根本上减少显存需求。 缺点：需要改模型，可能影响效果。

优先级 9：显存碎片优化

• 设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True；
• 避免频繁分配释放；
• 使用显存池；
• 减少不同大小 tensor 的交替分配。

优点：不改变计算，可能立刻见效。 缺点：收益取决于碎片情况。

优先级 10：Offload

• 把优化器状态放 CPU；
• 把参数放 CPU，用时再加载；
• 把激活值放 CPU。

优点：显存需求大幅下降。 缺点：速度明显变慢，受 PCIe 带宽限制。

---

四、按 OOM 位置的排查思路

1、模型加载时 OOM

• 模型参数太大；
• 解决：用 ZeRO-3、TP、量化加载、CPU Offload。

2、前向传播时 OOM

• 激活值太大；
• 解决：减小 batch、梯度检查点、混合精度。

3、反向传播时 OOM

• 激活值 + 梯度；
• 解决：梯度检查点、ZeRO-2、减小 batch。

4、优化器更新时 OOM

• 优化器状态太大；
• 解决：ZeRO-1、8-bit Adam、Adafactor、CPU Offload。

5、保存 checkpoint 时 OOM

• 显存汇聚；
• 解决：分片保存、CPU 保存、ZeRO-3。

---

五、排查和优化流程总结

第一步，确认 OOM 位置和显存组成。

第二步，按代价从低到高尝试：

• 1、减小 batch size；
• 2、开启混合精度；
• 3、梯度累积；
• 4、梯度检查点；
• 5、优化器显存优化；
• 6、ZeRO-1/2/3；
• 7、并行策略；
• 8、模型结构优化；
• 9、显存碎片优化；
• 10、Offload。

第三步，每次只改一个变量，观察效果。

第四步，如果还是 OOM，考虑组合使用多种手段。

第五步，记录配置和显存数据，形成经验。

---

六、一个具体例子

假设 7B 模型，单卡 80GB，FP16 + Adam，batch size 32，序列长度 2048。

估算：

• 模型状态：$16 \times 7\text{B} = 112\text{GB}$，单卡已经放不下；
• 激活值：约几十 GB。

解决方案：

• 先用 ZeRO-2，模型状态降到约 $2\Psi + 14\Psi/N$；
• 如果还不够，用 ZeRO-3；
• 同时开启梯度检查点，降低激活值；
• 如果还 OOM，减小 batch，用梯度累积；
• 最后考虑 TP + PP。

---

七、简单评价这道面试题

1、这道题的价值

• 这是大模型训练中最实战的一道题；
• 一道题能考察显存组成、优化手段、优先级判断和工程经验；
• 区分度好，初级可能只知道减小 batch，中级能说清 ZeRO 和梯度检查点，高级能按优先级系统排查；
• 贴近实战，每个训练大模型的人都会遇到 OOM。

2、这道题的不足

• 题目偏开放，如果没有实际经验，容易答得零散；
• 没有具体场景约束，比如“什么模型、多少卡、什么硬件”；
• 对显存碎片、Offload 等细节考察不够；
• 没有涉及 OOM 和训练稳定性的关系。

3、建议的追问方向

• 给定一个 70B 模型和 8 张 80GB 卡，怎么配置才能不 OOM？
• 梯度检查点和 ZeRO-3 一起用要注意什么？
• 显存碎片怎么排查和优化？
• CPU Offload 的代价有多大？
• 如何在不影响收敛的前提下降低显存？

4、总体评价

这是一道很好的大模型训练实战面试题，适合考察候选人对显存优化和工程调优的系统理解。它既能考察显存组成的理论，也能延伸到实际排查和优化。如果候选人能把 OOM 的定位方法、优化手段的优先级、各手段的代价和适用场景都讲清楚，基本可以判断其具备大模型训练显存优化的扎实基础。

---

---