3、诊断工具入门：nvidia-smi详解、gpustat使用、PyTorch Profiler基础
好，咱们开始聊诊断工具。说实话，排查显存爆掉的问题，就跟医生看病一样——你得先有听诊器、血压计，才能对症下药。nvidia-smi、gpustat、PyTorch Profiler，就是咱们这行的三大件。
我个人习惯是：先用nvidia-smi快速扫一眼，再用gpustat看个动态变化，最后用Profiler做精准定位。这套组合拳打下来，90%的显存问题都能找到根因。
3.1 nvidia-smi：最基础的显存监控
nvidia-smi是NVIDIA自带的命令行工具。你只要装了驱动，它就跟着装好了。我当年刚入行时，就靠这一条命令吃饭。
先看最常用的用法：
nvidia-smi输出大概长这样：
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.86 Driver Version: 470.86 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 45C P0 29W / 70W| 3246MiB / 15109MiB | 0% Default |
+-------------------------------+----------------------+----------------------+这里有几个关键字段，你得盯紧了：
| 字段 | 含义 | 我的经验 |
|---|---|---|
| Memory-Usage | 当前显存占用 / 总显存 | 如果接近上限，基本就是爆了 |
| GPU-Util | GPU计算单元利用率 | 低利用率+高显存占用，说明模型没跑起来 |
| Temp | 温度 | 超过80°C就要注意散热了 |
| Pwr:Usage/Cap | 当前功耗/最大功耗 | 功耗上不去，可能是CPU瓶颈 |
nvidia-smi -l 1 可以每秒刷新一次。我排查问题时，经常开两个终端：一个跑训练，一个跑这个命令实时盯着。
还有一个我常用的参数：
nvidia-smi --query-gpu=memory.used,memory.total --format=csv这个输出是纯文本格式，方便写脚本去解析。我在项目中就写过监控脚本，一旦显存占用超过90%，自动发告警。
3.2 gpustat：更友好的实时监控
gpustat是nvidia-smi的升级版。它用Python写的，安装很简单：
pip install gpustat用法也简单：
gpustat输出长这样：
GPU 0 [Tesla T4] 45°C, 29W | 3246/15109 MB | 0% | Default |
Processes:
12345 python 1234 MB
23456 python 2012 MB你看，它把每个进程的显存占用都列出来了。我曾经遇到一个情况：nvidia-smi显示显存用了80%，但我的模型只占了60%。一查gpustat才发现，有个同事的进程忘了关，偷偷占了20%。
gpustat还支持颜色高亮：
gpustat --color显存占用高的进程会标红，一眼就能看出来。
3.3 PyTorch Profiler：精准定位显存问题
前面两个工具只能告诉你「显存爆了」，但没法告诉你「为什么爆」。PyTorch Profiler就是干这个的。
先看一个最简单的用法：
import torch
from torch.profiler import profile, record_function, ProfilerActivity
with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("model_inference"):
output = model(input)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))输出会告诉你：每个操作占了多少显存、花了多少时间。我一般重点关注两个指标：
- cuda_memory_usage：显存占用最高的操作
- self_cuda_memory_usage：排除子操作后的自身显存占用
举个例子，我曾经发现一个卷积层占了80%的显存。仔细一看，是输入分辨率太大了。后来加了几个池化层，显存直接降了一半。
Profiler还能生成火焰图：
prof.export_chrome_trace("trace.json")这个文件可以用Chrome的chrome://tracing打开。火焰图能直观看到：哪个操作最耗时、哪个操作最占显存。
3.4 知识体系总览
下面这张图，是我自己总结的诊断工具使用流程。你照着这个思路走，排查效率能翻倍。
3.5 避坑指南
最后，分享几个我踩过的坑：
- 我曾经以为nvidia-smi显示的显存就是实际占用。 其实不是，它显示的是「已分配」的显存，有些显存虽然分配了但没在用。真正的占用要看
nvidia-smi --query-gpu=memory.used。 - gpustat默认只显示当前用户进程。 如果你用sudo跑，才能看到所有用户的进程。我一开始不知道，查了半天找不到那个「幽灵进程」。
- Profiler的record_shapes=True一定要开。 不开的话，你看不到每个张量的形状，定位问题就像盲人摸象。
嗯，工具就这些。说白了，诊断显存问题没那么玄乎——先看全局，再看进程，最后定位代码。这套流程走下来，基本没有搞不定的。