3、显存分析工具实战：nvidia-smi详解、torch.cuda.memory_summary()、memory_profiler库、PyTorch Profiler

做深度学习训练，最头疼的问题之一就是显存爆炸。

模型跑着跑着，突然给你来个 CUDA out of memory。这时候你怎么办？

我刚开始做大模型训练时，遇到这种错误就慌。后来发现，其实只要用好几个工具，显存问题基本都能定位。今天我就把这几个压箱底的工具分享给你。

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

nvidia-smi 是 NVIDIA 自带的命令行工具。说白了，它就是显卡的「体检报告」。

核心用法：

# 实时监控，每秒刷新一次
watch -n 1 nvidia-smi

# 查看更详细的显存信息
nvidia-smi --query-gpu=index,name,memory.total,memory.used,memory.free --format=csv

我个人习惯用 watch -n 1 nvidia-smi。为什么？因为训练时显存变化很快，每秒刷新才能看到真实情况。

输出结果里，有几个关键字段你要注意：

字段	含义	我的经验
Memory-Usage	当前显存使用量	如果接近 100%，说明快爆了
Volatile GPU-Util	GPU 计算单元利用率	这个值低但显存高，说明有瓶颈
Processes	占用 GPU 的进程	看看是不是有残留进程没杀掉

避坑指南：

我曾经遇到过一个问题：明明代码里调用了 torch.cuda.empty_cache()，但显存就是没释放。后来发现，是另一个 Python 进程还占着 GPU。用 nvidia-smi 一看，果然有个僵尸进程。杀掉就好了。

3.2 torch.cuda.memory_summary()：PyTorch 的显存放大镜

nvidia-smi 只能看到「谁用了多少显存」。但具体是哪个张量、哪个操作占的？它看不出来。

这时候就需要 torch.cuda.memory_summary() 了。

import torch

# 在训练的关键位置插入
print(torch.cuda.memory_summary())

# 或者只打印当前分配情况
print(torch.cuda.memory_allocated() / 1024**3, "GB")

这个函数会输出一份详细的报告，包括：

当前分配了多少显存
缓存池里有多少
每个张量的大小和地址
有没有内存碎片

嗯，这里要注意：memory_summary() 的输出很长。我建议只在怀疑有内存泄漏时用，平时别一直开着。

小技巧：

我习惯在训练循环的每个 epoch 结束时打印一次 memory_allocated()。如果数值一直在涨，说明有张量没释放。这时候就要检查是不是有 .detach() 没加，或者梯度没清零。

3.3 memory_profiler：Python 级别的内存分析

有时候显存问题不是 PyTorch 的锅，而是 Python 代码本身占用了太多内存。

memory_profiler 这个库，就是用来分析 Python 代码内存的。

# 安装
pip install memory_profiler

# 在函数上加装饰器
from memory_profiler import profile

@profile
def train_model():
    # 你的训练代码
    pass

运行后，它会输出每一行的内存变化。比如：

Line #    Mem usage    Increment   Line Contents
================================================
    10     50.2 MiB      0.0 MiB   def train_model():
    11     50.2 MiB      0.0 MiB       model = MyModel()
    12    150.2 MiB    100.0 MiB       data = load_data()  # 这里加载了数据
    13    250.2 MiB    100.0 MiB       output = model(data)  # 这里前向传播

你看，哪一行占了多少内存，一目了然。

注意：

这个工具会拖慢运行速度。我一般只在调试阶段用，正式训练时会关掉。

3.4 PyTorch Profiler：全链路性能分析

最后这个工具，是我最常用的。PyTorch Profiler 不仅能看显存，还能看算力、带宽、耗时。

import torch.profiler as profiler

with profiler.profile(
    activities=[
        profiler.ProfilerActivity.CPU,
        profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    profile_memory=True,
) as prof:
    # 你的训练代码
    output = model(input_data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

# 打印结果
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))

输出结果会告诉你：

每个操作消耗了多少显存
每个操作花了多少时间
有没有显存碎片
算力利用率怎么样

举个例子，我上次用 Profiler 发现，一个 attention 操作占了 40% 的显存。后来换成 Flash Attention，显存直接降了一半。

我的建议：

Profiler 的输出可以导出为 Chrome 的 trace 文件。用 chrome://tracing 打开，能看到非常直观的时间线。哪个操作慢、哪个占显存多，一眼就能看出来。

3.5 知识体系总览

这四个工具，各有各的用处。我画了张图，帮你理清它们的关系：

你看，这四个工具覆盖了从「整体监控」到「逐行分析」的完整链路。平时用 nvidia-smi 盯着，出问题了用 memory_summary() 和 memory_profiler 定位，最后用 Profiler 做深度优化。

说白了，显存优化不是玄学。只要工具用对了，问题都能找到根因。