4、诊断与监控工具：nvidia-smi内存使用查看、PyTorch memory_summary()、torch.cuda.memory_stats()、Nsight Systems内存轨迹分析
做显存优化，第一步不是动手改代码。
而是先搞清楚——显存到底被谁吃了？
我见过太多同学，一上来就调batch size、换模型结构，折腾半天发现根本问题不在那儿。所以这一章，咱们先把“诊断”这件事聊透。
4.1 nvidia-smi：最基础的“望闻问切”
nvidia-smi 是每个搞GPU的人都会用的命令。但说实话，大部分人只看一眼显存占用率就完事了。这远远不够。
关键指标解读：
- Memory-Usage：当前进程占用的显存。注意，这个值是“分配量”，不是“使用量”。PyTorch的缓存分配器会预占一大块，导致你看上去显存满了，实际很多是空闲缓存。
- Volatile GPU-Util：GPU计算单元利用率。如果这个值很低但显存爆了，说明你的模型在“占着茅坑不拉屎”。
- Processes：显示哪个进程在占用显存。我曾经排查过一个诡异问题，发现有个残留的Python进程占了8GB显存，kill掉之后一切正常。
我的习惯：在训练脚本里每隔几步打印一次 nvidia-smi 的关键字段。用 os.popen('nvidia-smi --query-gpu=memory.used --format=csv') 直接拿到数值，比肉眼盯着终端靠谱得多。
4.2 PyTorch memory_summary()：你的显存“体检报告”
nvidia-smi 只能看到“总量”，看不到内部细节。这时候就需要 PyTorch 自带的 memory_summary() 了。
import torch
# 在训练循环中插入
print(torch.cuda.memory_summary())
# 或者更轻量的
print(torch.cuda.memory_allocated()) # 实际分配
print(torch.cuda.memory_reserved()) # 缓存预留这个函数会输出一份详细的报告，包括：
- 当前分配了多少个张量
- 每个张量的大小和类型
- 缓存池的使用情况
- 碎片化程度（通过“Large Pool”和“Small Pool”的对比）
我曾经踩过的坑：memory_summary() 的输出非常长，尤其是在大模型上。我建议只在关键节点调用，比如每个epoch结束，或者显存异常时。否则打印本身就会拖慢训练速度。
4.3 torch.cuda.memory_stats()：程序员的“显微镜”
memory_summary() 是给人看的。memory_stats() 是给程序看的。
它返回一个字典，里面全是数字。我常用的几个字段：
stats = torch.cuda.memory_stats()
print(stats['allocated_bytes.all.current']) # 当前分配字节数
print(stats['reserved_bytes.all.current']) # 当前预留字节数
print(stats['active_bytes.all.current']) # 活跃字节数（分配且未释放）
print(stats['num_alloc_retries']) # 分配重试次数（碎片化的重要指标）为什么 num_alloc_retries 很重要？
当PyTorch尝试分配显存但失败时，它会触发一次“垃圾回收+碎片整理”。如果这个数字很大，说明你的显存碎片化已经很严重了。我一般会在训练脚本里监控这个值，一旦超过某个阈值就报警。
我的做法：写一个简单的装饰器，在每个step结束后自动记录这些统计值，写入日志文件。这样事后分析时，可以精确定位到哪个step出了问题。
4.4 Nsight Systems：可视化“时间线”
前面几个工具都是“快照”，只能看到某一时刻的状态。但显存问题往往是动态的——某个操作突然分配了大量显存，然后又释放了。这时候就需要 Nsight Systems 了。
Nsight Systems 可以记录整个训练过程的内存分配和释放事件，生成一个时间线图。你能看到：
- 每个CUDA kernel的启动时间
- 每个张量的分配和释放时刻
- CPU和GPU之间的数据传输
- 显存使用量的变化曲线
使用方式很简单：
# 命令行
nsys profile -o trace_output python train.py
# 或者用PyTorch的profiler
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CUDA],
profile_memory=True,
record_shapes=True
) as prof:
train_one_epoch()
print(prof.key_averages().table(sort_by="self_cuda_memory_usage"))一个真实案例：我之前优化一个Transformer模型，用nvidia-smi看显存占用一直稳定在80%左右，但偶尔会突然跳到95%然后OOM。用Nsight Systems一看，发现是某个attention计算时，临时创建了一个巨大的中间张量，用完就释放了。但因为释放不及时，导致下一个大张量分配时撞上了碎片化区域，触发了OOM。
解决方案很简单：在那个操作前后手动调用了 torch.cuda.empty_cache()，问题就解决了。
4.5 工具选择指南
| 工具 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|
| nvidia-smi | 快速查看整体显存占用 | 零依赖，实时 | 看不到内部细节 |
| memory_summary() | 详细分析当前显存状态 | 信息全面，包含碎片化指标 | 输出冗长，影响性能 |
| memory_stats() | 自动化监控和告警 | 可编程，轻量 | 需要自己解析数据 |
| Nsight Systems | 定位动态内存问题 | 可视化时间线，精准定位 | 需要安装，有一定学习成本 |
我的建议：日常开发用 nvidia-smi + memory_stats() 就够了。遇到疑难杂症再上 Nsight Systems。别一上来就开大炮打蚊子。
4.6 本章知识体系
下面这张图总结了四个工具的关系和使用场景：
嗯，这一章的内容就到这里。工具是死的，人是活的。关键是要养成“先诊断、后优化”的习惯。下一章咱们会深入讲碎片化的成因，到时候这些工具就能派上大用场了。