2、显存监控：nvidia-smi详解、PyTorch显存监控工具、自定义显存监控脚本

显存监控这事儿，说白了就是给模型训练装个「仪表盘」。你想想看，开车没有仪表盘，油量耗尽了你都不知道。我见过太多同学，模型跑到一半突然OOM（显存溢出），然后一脸懵地问我：「老师，我代码明明没问题啊？」

嗯，问题往往不在代码逻辑，而在显存管理。这一章，我就带你把这套「仪表盘」装好、看懂、用熟。

2.1 nvidia-smi：最基础的显存监控工具

nvidia-smi 是 NVIDIA 官方提供的命令行工具。我个人习惯，每次跑模型前先敲一遍这个命令，看看当前显存状态。

2.1.1 基本用法

nvidia-smi

输出结果里，你需要重点关注这几列：

列名	含义	我的经验
Memory-Usage	显存使用量 / 总显存	比如 4096MiB / 8192MiB，就是用了4GB
Volatile GPU-Util	GPU计算单元利用率	这个值低不代表显存没压力
Processes	占用GPU的进程列表	看PID可以找到是哪个程序在吃显存

2.1.2 实时监控模式

watch -n 1 nvidia-smi

这个命令每1秒刷新一次。我跑大模型训练时，经常开两个终端：一个跑训练，一个跑这个监控。这样能第一时间发现显存异常增长。

2.1.3 只看显存信息

nvidia-smi --query-gpu=memory.used,memory.total --format=csv

输出类似：

memory.used [MiB], memory.total [MiB]
4096, 8192

这个格式很适合写脚本去解析。我后面会讲怎么用Python自动抓取这些数据。

2.2 PyTorch显存监控工具

nvidia-smi 能看到全局显存，但看不到PyTorch内部细节。比如，某个张量到底占了多少显存？有没有显存泄漏？这时候就需要PyTorch自带的工具了。

2.2.1 torch.cuda.memory_allocated()

这个函数返回当前已分配的显存（字节数）。

import torch

print(f"已分配显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存显存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

这里有个坑：memory_allocated 是实际使用的，memory_reserved 是PyTorch向GPU申请的缓存。我刚开始做优化时，发现 memory_reserved 比 memory_allocated 大很多，还以为有泄漏。后来才知道，这是PyTorch的缓存机制，为了加速后续分配。

2.2.2 torch.cuda.memory_summary()

这个函数会输出一份详细的显存报告，包括每个张量的分配情况。

print(torch.cuda.memory_summary())

输出内容很长，但有几个关键信息：

• 当前分配的张量数量
• 最大分配过的显存（峰值）
• 每个张量的地址和大小

2.2.3 显存快照与可视化

torch.cuda.memory._dump_snapshot("snapshot.pkl")

这个功能可以把当前显存状态保存下来，然后用Chrome浏览器打开 chrome://tracing 加载这个文件，就能看到一张显存分配的时间线图。

说实话，这个功能我平时用得不多，但遇到疑难杂症时特别好使。有一次我怀疑某个中间变量没释放，就是靠这个快照找到的。

2.3 自定义显存监控脚本

工具再好，也得整合起来用。我习惯写一个脚本，把nvidia-smi和PyTorch的监控信息汇总到一起，实时输出。

2.3.1 一个简单的监控脚本

import torch
import subprocess
import time

def get_gpu_info():
    """获取nvidia-smi的显存信息"""
    result = subprocess.run(
        ['nvidia-smi', '--query-gpu=memory.used,memory.total',
         '--format=csv,nounits,noheader'],
        capture_output=True, text=True
    )
    used, total = result.stdout.strip().split(', ')
    return int(used), int(total)

def monitor_memory(interval=2):
    """每interval秒打印一次显存状态"""
    while True:
        # nvidia-smi信息
        used, total = get_gpu_info()
        # PyTorch信息
        allocated = torch.cuda.memory_allocated() / 1024**2
        reserved = torch.cuda.memory_reserved() / 1024**2
        
        print(f"[{time.strftime('%H:%M:%S')}] "
              f"全局: {used}/{total} MB | "
              f"PyTorch已分配: {allocated:.0f} MB | "
              f"缓存: {reserved:.0f} MB")
        
        time.sleep(interval)

if __name__ == "__main__":
    monitor_memory()

这个脚本我一直在用。跑训练时，开一个终端跑这个监控，一目了然。

2.3.2 记录显存峰值

很多时候，我们不光要看当前显存，还要知道训练过程中显存最高到过多少。PyTorch提供了记录峰值的功能：

torch.cuda.reset_peak_memory_stats()
# 跑你的训练代码
peak = torch.cuda.max_memory_allocated() / 1024**2
print(f"显存峰值: {peak:.2f} MB")

2.4 本章知识体系

下面这张图，是我对显存监控工具的整体理解。你可以把它当作一个「监控工具箱」：

好了，这一章的内容就到这里。显存监控是优化的第一步，只有「看得清」，才能「改得准」。下一章，我们会聊怎么用这些监控工具去定位具体的显存瓶颈。

想要更多资料，联系可可