压测工具选型：nvidia-smi、nvtop、PyTorch Profiler、Nsight Systems对比

做显存瓶颈定位，第一步就是选对工具。

说实话，我见过不少同学一上来就开 PyTorch Profiler，结果发现显存占用高得离谱，但根本不知道是哪块代码在吃显存。也有同学盯着 nvidia-smi 看半天，发现显存快满了，但模型还在正常跑，完全不知道瓶颈在哪。

嗯，工具选不对，后面全是白费功夫。今天我就把这四个主流工具掰开揉碎讲清楚。

1. nvidia-smi：最基础的显存监控

这是 NVIDIA 官方自带的命令行工具，几乎每台 GPU 机器上都有。我个人习惯把它当作「第一道防线」——先看看显存是不是真的满了。

# 实时监控显存
watch -n 1 nvidia-smi

# 只看显存使用
nvidia-smi --query-gpu=memory.used,memory.total --format=csv

优点：零安装、零配置，看一眼就知道显存总量和当前使用量。

缺点：只能看到「谁在用显存」，看不到「为什么用」。比如你发现显存占了 20GB，但不知道是模型参数、梯度、还是中间激活值占的。

2. nvtop：可视化显存监控

nvtop 是 GPU 版的 htop，用起来很直观。我一般在需要「持续观察显存变化趋势」时用它。

# 安装（Ubuntu）
sudo apt install nvtop

# 启动
nvtop

优点：图形化界面，能看到每个 GPU 的显存、温度、功耗、进程列表。按 p 键还能排序。

缺点：和 nvidia-smi 一样，只能看到「谁在用」，看不到「为什么用」。而且它需要额外安装，有些生产环境不让装。

3. PyTorch Profiler：代码级别的显存分析

这才是真正能定位「哪行代码吃了显存」的工具。PyTorch 自带的 profiler 可以记录每个操作的显存分配和释放。

import torch
from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True
) as prof:
    # 你的模型推理代码
    output = model(input_tensor)

print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage",
    row_limit=20
))

输出示例：

-------------------------  ------------  ------------  ------------  ------------
Name                     Self CPU %     Self CPU      CUDA Mem      Self CUDA Mem
-------------------------  ------------  ------------  ------------  ------------
aten::linear             12.34%        5.2ms         2.1 GB        1.8 GB
aten::matmul             8.56%         3.6ms         1.5 GB        1.2 GB
aten::embedding          5.21%         2.1ms         0.8 GB        0.6 GB
-------------------------  ------------  ------------  ------------  ------------

优点：精确到每个 PyTorch 操作的显存分配，能告诉你「哪个算子吃了最多显存」。

缺点：只能分析 PyTorch 代码，对 CUDA 内核级别的显存分配无能为力。而且 profiling 本身会引入额外开销，生产环境慎用。

4. Nsight Systems：系统级显存分析

这是 NVIDIA 官方的高级分析工具，能看整个系统的 GPU 活动，包括 CUDA 内核、内存拷贝、显存分配等。我一般在「PyTorch Profiler 不够用」时用它。

# 命令行启动
nsys profile -o trace_output python train.py

# 打开 GUI 查看
nsys-ui trace_output.qdrep

优点：能看到 CUDA 内核级别的显存分配，包括 cudaMalloc、cudaFree 的调用栈。还能分析 CPU 和 GPU 之间的数据传输。

缺点：学习曲线陡峭，GUI 界面信息量太大，新手容易迷失。而且 profiling 开销比 PyTorch Profiler 大。

5. 工具对比总结

6. 我的选型建议

说白了，没有万能工具。我一般按这个流程来：

1. 第一步：用 nvidia-smi 或 nvtop 确认显存确实打满了。如果没打满，那瓶颈可能在别处（比如计算或 I/O）。
2. 第二步：用 PyTorch Profiler 找到「哪个算子吃显存最多」。这一步能解决 80% 的问题。
3. 第三步：如果 PyTorch Profiler 不够用（比如显存泄漏是 CUDA 内核级别的），再上 Nsight Systems。

7. 知识体系图

这张图把四个工具的定位和关系说清楚了。你想想看，从最粗的进程级监控，到最细的 CUDA 内核级分析，每个工具都有自己的生态位。选对工具，定位显存瓶颈就成功了一半。