4. 显存监控与调试工具：nvidia-smi、torch.cuda.MemoryStats、Nsight Systems 实操
显存监控这件事，说白了就是「你得知道你的钱（显存）花哪儿了」。很多同学训练跑着跑着就 OOM，然后一脸懵地跑来问我：「老师，我显存去哪了？」其实不是你代码写错了，是你还没学会查账。
我个人习惯，是在训练脚本启动的同时，至少开三个终端：一个跑训练，一个跑 nvidia-smi -l 1，还有一个留着随时敲 torch.cuda.memory_summary()。这不是矫情，是吃过亏之后的肌肉记忆。
4.1 nvidia-smi：最硬的底层快照
nvidia-smi 是所有监控的起点。它直接问驱动要数据，绕过了 PyTorch 的缓存层，所以看到的数字是「物理显存占用」。但要注意——它看到的是进程级总量，不是张量级明细。
# 每 1 秒刷新一次
nvidia-smi -l 1
# 只显示指定 GPU，适合多卡场景
nvidia-smi -i 0,1 -l 1
# 显示 GPU 进程详情，能看到进程 PID 和显存占用
nvidia-smi pmon -i 0
# 一次性查看可用的 GPU 编号和空闲显存
nvidia-smi --query-gpu=index,memory.free --format=csv我有一次在 8 卡机器上训模型，总显存被吃掉 60GB，但看每个进程只占了 4GB，怎么都对不上。后来才发现是 NCCL 的缓存池和 PyTorch 的缓存层重叠统计了——这就是 nvidia-smi 的局限性，它只告诉你「物理上用了多少」，不告诉你「逻辑上你的张量占了多少」。
| 字段 | 含义 | 我常用的判断 |
|---|---|---|
| Volatile GPU-Util | GPU 核心利用率（%） | 低于 80% 时检查数据加载是否成了瓶颈 |
| Memory-Usage | 物理显存占用（MiB） | 接近 90% 就要警惕 OOM |
| Processes / GPU Memory | 每个进程占用的显存 | 超过单进程上限时考虑梯度 checkpoint |
nvidia-smi 看显存占用 0%，以为模型没加载进去，其实是 PyTorch 的缓存分配器还没真正向驱动申请物理页——PyTorch 是 lazy allocation。这时候需要手动跑一步 forward，再回头看才是真实占用。
4.2 torch.cuda.MemoryStats：PyTorch 内部账本
比 nvidia-smi 更细的，是 PyTorch 自己管的账。说白了，PyTorch 有自己的显存分配器。它把从驱动拿到的一大块物理显存切成小块分配给张量，用完不还回去，放回缓存池里等下次复用。这个机制叫 CUDA caching allocator。
所以 nvidia-smi 看到的是「PyTorch 从驱动那拿了多少」，而 torch.cuda.memory_stats() 看到的是「PyTorch 内部怎么分的」。
import torch
# 1. 快速看当前分配和保留量
print(torch.cuda.memory_allocated()) # 实际被张量占用的字节数
print(torch.cuda.memory_reserved()) # 缓存分配器从驱动拿到的总字节数
# 2. 打印详细的分配器段信息
print(torch.cuda.memory_stats())
# 3. 人类可读的摘要，推荐放在每个 epoch 结尾
print(torch.cuda.memory_summary())memory_summary() 的输出会把所有内存池（small_pool、large_pool）的细节列出来。你注意看 Large Pool 这一栏——如果它里面有很多 used 但 free 很少，说明你的张量生命周期太长了，该及时 del 或者用 torch.no_grad() 包住推理部分。
# 估算每个 batch 的显存增量
prev = torch.cuda.memory_allocated()
loss = model(batch) # 你的前向 + 反向
delta = torch.cuda.memory_allocated() - prev
print(f"当前 batch 显存增量: {delta / 1024**2:.1f} MB")你想想看，如果这个增量突然变大，十有八九是某个中间张量没有被释放——最常见的是把 loss.item() 的 tensor 留在计算图里了。
4.3 用 MemorySnail 做时间线追踪
等等，我要先推荐一个我觉得特别好用的小工具——MemorySnail。它是 Meta 开源的内存调试脚本，本质是定时轮询 torch.cuda.memory_stats()，然后把数据画成时间线图。
用法很简单，把它作为一个后台线程启动，训练跑完后会生成一个 HTML，里面有一张按时间轴的显存分配拆解图。
你会清晰地看到：哪一秒梯度 checkpoint 帮你省了 2GB，哪一秒 Adam 优化器状态暴涨了一截。
# 安装
pip install memory_snails
# 在你的训练脚本里启动
from memory_snails import SnailTracer
tracer = SnailTracer(interval=0.5) # 每 0.5 秒采样一次
tracer.start()
# 你的训练循环
for epoch in range(num_epochs):
train_one_epoch()
tracer.stop()
tracer.save("memory_timeline.html")4.4 Nsight Systems：性能与内存的合体诊断
前面几个工具只盯着显存。但如果要同时看「显存怎么涨的」和「GPU 计算单元为什么在摸鱼」，那就得上 Nsight Systems 了。
Nsight Systems 是 NVIDIA 官方的全系统性能分析器。它用跟踪采样的方式，把 CPU 端和 GPU 端的事件按时间轴对齐。你能看到：
- 每个 CUDA kernel 的启动和结束时间
- 数据从 CPU 传到 GPU 的 copy 事件
- 每次 cudaMalloc 和 cudaFree 的调用点
- PyTorch 算子的调用堆栈（需要符号表支持）
我最常用的采集命令是这样：
# 先清空缓存，然后跑 100 步训练
nsys profile --stats=true -o my_trace \
-t cuda,nvtx,osrt \
--force-overwrite true \
python train.py --max_steps 100跑完之后会生成一个 .nsys-rep 文件。用 Nsight Systems GUI（图形界面）打开它，你会看到类似 Gantt 图的时间轴。
| 你会在时间轴上看到什么 | 对应的问题 |
|---|---|
| 大段的空白没有 kernel 执行 | CPU 数据加载或预处理是瓶颈 |
| cudaMalloc 调用频繁且耗时 | 张量生命周期过短，考虑用缓存池或复用 tensor |
| 小 kernel 密集启动，每个都很短 | 算子融合不够，kernel launch overhead 太高 |
| H2D/D2H memcpy 耗时占比 > 5% | 数据加载未做 pinned memory 或者未预取 |
我建议你用 Nsight Systems 配合 torch.cuda.nvtx.range_push 加标记，这样一眼就能把时间轴上的内存分配和具体代码段对应上。
# 在训练循环里加上 NVTX 标记
import torch.cuda.nvtx as nvtx
for batch in dataloader:
nvtx.range_push("forward")
out = model(batch)
nvtx.range_pop()
nvtx.range_push("backward")
loss.backward()
nvtx.range_pop()
nvtx.range_push("optimizer_step")
optimizer.step()
nvtx.range_pop()加了标记之后，在 Nsight Systems 的时间轴上，每段代码的起止和对应的 kernel、内存操作就一目了然了——这对于定位「显存泄漏发生在哪一步」特别管用。
4.5 三款工具的定位速查
说句大实话：日常调试我 80% 的时间只开 nvidia-smi 和 torch.cuda.memory_summary()。只有在遇到「内存明明有富余但跑着跑着就崩了」或者「GPU 利用率怎么都提不上去」这种疑难杂症时，我才会起 Nsight Systems。记住一个原则——先看日志，再上探测器。不要一上来就搞得跟军火库似的，你累我也累。