2. 显存监控工具

显存监控这事儿，说实话，是很多AI训练工程师容易忽略的一环。我刚开始做大规模训练时，也踩过这个坑——模型跑着跑着突然OOM，然后一脸懵地开始猜到底哪一步爆了显存。后来我学乖了，先把监控工具玩明白，再动手调优。

这一章，我就带你把这几个核心工具捋一遍。从最基础的nvidia-smi，到可视化神器nvtop，再到PyTorch自带的Profiler，最后我们自己动手写个监控脚本。嗯，一步步来。

2.1 nvidia-smi：你每天都在用，但可能没用好

nvidia-smi是NVIDIA官方提供的命令行工具，几乎每个用GPU的人都会敲。但说实话，大部分人只用了它20%的功能。

基础用法

# 最常用的查看GPU状态
nvidia-smi

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

输出里你会看到显存使用量、GPU利用率、温度、功耗等信息。我个人习惯重点关注两个指标：显存使用量和GPU利用率。前者告诉你显存够不够，后者告诉你计算单元有没有在干活。

进阶用法：查询进程级显存占用

# 查看每个进程的显存占用
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv

# 输出示例
pid, process_name, used_memory [MiB]
12345, python, 10240
12346, python, 8192

我在项目中遇到过一个问题：明明只跑了一个模型，显存却占了30多GB。用这个命令一查，发现有个残留的Python进程占着20GB没释放。嗯，这种坑很常见。

2.2 nvtop：让显存监控可视化

nvtop是GPU版的htop，界面友好，信息一目了然。我个人非常推荐，尤其是调试时盯着终端看数字变化太枯燥了。

安装

# Ubuntu/Debian
sudo apt install nvtop

# 或者从源码编译
git clone https://github.com/Syllo/nvtop.git
mkdir -p nvtop/build && cd nvtop/build
cmake ..
make
sudo make install

使用

nvtop

启动后你会看到每个GPU的实时状态，包括显存使用曲线、进程列表、温度等。按F2可以设置刷新频率，我一般设成500ms。

2.3 PyTorch Profiler：精准定位显存瓶颈

前面两个工具只能看到全局显存占用，但具体是哪一层、哪个操作吃了显存，它们就无能为力了。这时候就需要PyTorch自带的Profiler出马。

基本用法

import torch
from torch.profiler import profile, record_function, ProfilerActivity

# 定义一个简单的模型
model = torch.nn.Linear(1024, 1024).cuda()
input_tensor = torch.randn(1024, 1024).cuda()

with profile(
    activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True
) as prof:
    with record_function("model_inference"):
        output = model(input_tensor)

# 打印显存占用最高的操作
print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))

输出会告诉你每个操作（比如矩阵乘法、激活函数）占了多少显存。我在调优一个Transformer模型时，发现softmax操作居然占了20%的显存，后来换成flash_attention直接降了一半。

进阶：追踪训练循环

from torch.profiler import profile, ProfilerActivity

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())

with profile(
    activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
    profile_memory=True,
    record_shapes=True,
    with_stack=True
) as prof:
    for step in range(10):
        optimizer.zero_grad()
        loss = model(input_tensor)
        loss.backward()
        optimizer.step()
        prof.step()  # 标记每一步

# 按显存占用排序
print(prof.key_averages(group_by_input_shape=True).table(
    sort_by="self_cuda_memory_usage", row_limit=20))

2.4 自定义显存监控脚本

有时候现成的工具不够灵活，比如你想在训练过程中自动记录显存峰值、或者当显存超过阈值时自动报警。这时候就得自己写脚本了。

方案一：基于nvidia-smi的Python封装

import subprocess
import re
import time

def get_gpu_memory():
    """获取所有GPU的显存使用情况"""
    result = subprocess.run(
        ['nvidia-smi', '--query-gpu=index,memory.used,memory.total',
         '--format=csv,noheader,nounits'],
        capture_output=True, text=True
    )
    lines = result.stdout.strip().split('\n')
    gpu_info = []
    for line in lines:
        parts = line.split(', ')
        gpu_info.append({
            'index': int(parts[0]),
            'used_mb': int(parts[1]),
            'total_mb': int(parts[2])
        })
    return gpu_info

def monitor_memory(interval=1, threshold=0.9):
    """监控显存，超过阈值报警"""
    print("开始监控显存...")
    while True:
        gpu_info = get_gpu_memory()
        for gpu in gpu_info:
            usage_ratio = gpu['used_mb'] / gpu['total_mb']
            if usage_ratio > threshold:
                print(f"[警告] GPU {gpu['index']} 显存使用率: {usage_ratio:.1%}")
            else:
                print(f"GPU {gpu['index']}: {gpu['used_mb']}/{gpu['total_mb']} MB ({usage_ratio:.1%})")
        time.sleep(interval)

if __name__ == '__main__':
    monitor_memory()

方案二：集成到训练代码中

import torch
import time

class MemoryMonitor:
    def __init__(self, device_id=0):
        self.device_id = device_id
        self.peak_memory = 0

    def step(self):
        """记录当前显存，更新峰值"""
        current = torch.cuda.memory_allocated(self.device_id)
        self.peak_memory = max(self.peak_memory, current)
        return current

    def reset_peak(self):
        self.peak_memory = 0

# 在训练循环中使用
monitor = MemoryMonitor()
for epoch in range(10):
    monitor.reset_peak()
    for batch in dataloader:
        # 前向传播
        loss = model(batch)
        # 记录显存
        current_mem = monitor.step()
        if current_mem > 0.9 * torch.cuda.get_device_properties(0).total_memory:
            print(f"显存告警！当前使用: {current_mem / 1024**3:.2f} GB")
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch} 峰值显存: {monitor.peak_memory / 1024**3:.2f} GB")

2.5 知识体系总览

下面这张图把本章的核心内容串起来了，方便你理解各个工具之间的关系和适用场景。

这张图展示了四个工具的定位：nvidia-smi负责全局概览，nvtop提供可视化体验，PyTorch Profiler深入操作级别，自定义脚本则实现自动化监控。实际项目中，我通常组合使用——先用nvidia-smi快速扫一眼，发现异常后用Profiler精确定位，最后写个脚本固化监控流程。

好了，这一章的内容就到这里。工具是死的，人是活的。多动手试试，找到最适合自己工作流的监控方案。