3、显存溢出定位方法：使用nvidia-smi、PyTorch显存分析工具、日志监控等手段定位溢出点
显存溢出（CUDA Out of Memory，简称 OOM）是深度学习训练中最常见的故障之一。它往往不是单一原因造成的，而是模型参数、激活值、优化器状态、临时缓冲区以及碎片化等多重因素叠加的结果。因此，定位溢出点的关键不在于“看到报错就调小 batch size”，而在于建立一套从硬件层、框架层到应用层的分层观测体系，逐步缩小问题范围，找到真正吃掉显存的那一部分。
一、核心原理：显存被谁占用了
在 PyTorch 训练中，显存占用大致可以分为以下几类：模型参数（Parameters）、梯度（Gradients）、优化器状态（Optimizer States）、前向激活值（Activations）、临时工作区（Workspace）以及 CUDA 上下文与缓存分配器保留的显存。不同训练阶段，这些部分的占比会动态变化。例如，前向传播时激活值快速增长，反向传播时梯度与激活值同时存在，优化器更新时还会额外产生状态张量。
定位 OOM 的本质，就是判断“在哪个阶段、哪一类张量、以什么规模”导致了显存峰值超过物理上限。因此，我们需要三类工具协同工作：nvidia-smi 用于观察 GPU 整体显存与进程占用；PyTorch 显存分析工具用于细粒度追踪张量分配与释放；日志监控用于记录训练过程中的显存变化趋势，捕捉溢出前的异常信号。
二、方法流程：三步定位法
第一步：用 nvidia-smi 建立宏观视图。当训练报出 OOM 时，首先执行以下命令，确认 GPU 是否真的被占满，以及是否有其他进程残留。
nvidia-smi
nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv如果发现显存被多个进程占用，需要先清理僵尸进程。如果只有一个训练进程，但显存占用接近上限，则说明是训练本身导致的溢出。此时可以周期性采样，观察显存增长曲线：
while true; do
nvidia-smi --query-gpu=timestamp,memory.used --format=csv,noheader
sleep 2
done第二步：用 PyTorch 工具定位张量级分配。PyTorch 提供了丰富的显存查询接口，可以在训练循环中插入监控代码：
import torch
def log_memory(step):
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
max_allocated = torch.cuda.max_memory_allocated() / 1024**3
print(f"step={step} allocated={allocated:.2f}G "
f"reserved={reserved:.2f}G max={max_allocated:.2f}G")
# 在每个 step 后调用
log_memory(step)其中，memory_allocated 表示当前张量实际占用的显存，memory_reserved 表示 PyTorch 缓存分配器向 CUDA 申请的总显存，max_memory_allocated 表示历史峰值。如果 reserved 远大于 allocated，说明存在碎片化或缓存未释放；如果 max 在某个 step 突然跳升，则说明该 step 产生了大张量。
更细粒度的分析可以使用 torch.cuda.memory_summary()，它会输出分配器状态、碎片比例、各段大小等详细信息。对于复杂模型，还可以使用 torch.cuda.memory._record_memory_history() 记录分配历史，再通过快照分析具体是哪个操作分配了最大块显存。
第三步：用日志监控回溯溢出点。在训练脚本中加入结构化日志，记录每个 step 的 loss、学习率、序列长度、batch size 以及显存峰值。当 OOM 发生时，最后一条日志往往就是溢出前的现场。例如：
[step 1280] loss=2.341 lr=1e-4 seq_len=1024 batch=8
[step 1281] loss=2.337 lr=1e-4 seq_len=1024 batch=8
[step 1282] loss=2.352 lr=1e-4 seq_len=2048 batch=8
CUDA out of memory. Tried to allocate 2.00 GiB从日志可以清晰看到，step 1282 的序列长度从 1024 跳到了 2048，激活值规模翻倍，直接导致溢出。这种“日志 + 显存”联合分析，比单纯看报错信息有效得多。
三、案例场景：一次典型的 OOM 定位
某团队在训练一个 7B 参数模型时，使用 A100 40GB，batch size 为 4，序列长度 2048，开启 Adam 优化器。训练到第 300 步左右频繁 OOM。按照三步法定位：
- nvidia-smi 显示显存占用 39.2GB / 40GB，仅一个训练进程，排除多进程干扰。
- PyTorch 监控 显示 allocated 约 28GB，reserved 约 38GB，碎片化严重；max_memory_allocated 在 OOM 前一步达到 37.5GB。
- 日志 显示 OOM 前一步的 loss 正常，但梯度范数突然增大，触发了梯度裁剪前的临时大张量分配。
最终定位到两个问题：一是 Adam 优化器状态占用约 14GB（7B 参数 × 2 个状态 × 4 字节），二是注意力计算中的临时缓冲区在长序列下膨胀。解决方案是引入 ZeRO-2 分片优化器状态，并启用梯度检查点，显存峰值降至 26GB，训练稳定运行。
四、误区与复盘
误区一：只看 nvidia-smi 就下结论。nvidia-smi 显示的是 GPU 整体显存，无法区分是参数、激活还是缓存。必须结合 PyTorch 工具才能定位到具体张量。
误区二：忽略 reserved 与 allocated 的差异。reserved 高但 allocated 低，说明缓存分配器保留了大量空闲块，可能是碎片化，也可能是正常缓存。此时应关注 memory_reserved 与 memory_allocated 的比值。
误区三：OOM 后只调小 batch size。调小 batch size 可能掩盖问题，但如果是序列长度突增或优化器状态过大，调 batch 效果有限。应先定位再调参。
复盘建议：建立显存基线档案，记录不同配置下的 allocated、reserved、max 峰值；在训练脚本中固定加入显存日志；对长序列任务提前做显存估算，避免运行到一半才 OOM。
| 工具/手段 | 观测对象 | 适用阶段 | 关键指标 |
|---|---|---|---|
| nvidia-smi | GPU 整体显存与进程 | 宏观排查 | memory.used / free |
| torch.cuda.memory_allocated | 张量实际占用 | 训练循环内 | allocated |
| torch.cuda.memory_reserved | 缓存分配器保留量 | 训练循环内 | reserved |
| max_memory_allocated | 历史峰值 | 训练循环内 | max |
| memory_summary | 分配器详细状态 | OOM 后分析 | 碎片比例、段大小 |
| 结构化日志 | step 级上下文 | 全程 | seq_len、batch、loss |
掌握这三类工具的组合使用，就能把“显存溢出”从一个模糊的报错，转化为可追踪、可验证、可优化的具体问题。定位越精准，后续的 ZeRO 优化和调参才越有针对性。