2. 显存监控：nvidia-smi、torch.cuda、PyTorch Profiler 与自定义工具

显存监控这事儿，说实话，是每个做深度学习的人绕不开的坎儿。我刚开始做大模型训练那会儿，经常遇到「CUDA out of memory」的报错，然后一脸懵地到处找原因。后来慢慢摸索出一套监控体系，今天就跟大家聊聊。

显存监控说白了就两件事：看用了多少 和 看谁在用。前者帮你判断是否超限，后者帮你定位内存泄漏或显存碎片。下面我按工具逐个拆解。

2.1 nvidia-smi：最基础的显存快照

这个命令大家肯定不陌生。我习惯在训练脚本里每隔几轮就调用一次，记录显存峰值。

# 终端直接运行
nvidia-smi

# 只看显存信息，简洁模式
nvidia-smi --query-gpu=memory.used,memory.free,memory.total --format=csv

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

输出大概长这样：

GPU	Memory-Usage	Volatile GPU-Util
0	12345MiB / 24576MiB	85%
1	6789MiB / 24576MiB	32%

小技巧： 我经常用 nvidia-smi -l 1 来每秒刷新一次，但注意这个命令会占用终端。更好的做法是写个后台脚本，把日志重定向到文件里。

注意： nvidia-smi 显示的显存占用是「进程级别」的，不是「张量级别」。也就是说，它告诉你某个进程占了多少显存，但不知道里面是模型参数、梯度还是优化器状态。想细粒度分析，得用下面的工具。

2.2 torch.cuda：Python 里的实时监控

在代码里直接查显存，比切出去看 nvidia-smi 方便多了。PyTorch 提供了一整套 API。

import torch

# 当前设备显存使用情况（单位：字节）
print(f"已分配显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存显存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
print(f"最大已分配: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")

# 重置峰值统计
torch.cuda.reset_peak_memory_stats()

# 查看指定设备
print(torch.cuda.memory_summary(device=0))

我个人习惯在训练循环的每个 step 结束时，打印一次显存峰值。这样能直观看到哪个 batch 显存暴涨。

  核心概念区分：
  memory_allocated：实际被张量占用的显存
memory_reserved：PyTorch 向 CUDA 申请的缓存池大小（通常比 allocated 大）
两者差值就是「预留但未使用」的显存，这部分可以释放给其他进程

为什么会关注这个差值？我曾经遇到一个情况：显存明明还有 2GB 空闲，但新张量就是分配失败。后来发现是 PyTorch 的缓存机制导致的——它预留了一大块显存，但没实际用，其他进程也拿不到。这时候调用 torch.cuda.empty_cache() 可以释放缓存。

2.3 PyTorch Profiler：深度分析利器

如果你想知道「哪个操作吃了最多显存」，那就得上 Profiler 了。这工具不仅能看显存，还能看算力、带宽、耗时。

import torch.profiler as profiler

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

# 打印显存占用排名
print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage",
    row_limit=10
))

输出会告诉你：aten::mm 占了多少显存，native_layer_norm 占了多少。我曾在一次调优中发现，某个自定义的 attention 实现比官方版本多吃了 30% 的显存，就是靠 Profiler 抓出来的。

避坑指南： 我曾经开着 Profiler 跑完整训练流程，结果显存直接翻倍了——因为 Profiler 本身也要记录数据。建议只 profile 几个 step，或者用 schedule 参数做周期性采样。

2.4 自定义显存监控工具

有时候现成工具不够用，比如你想监控「模型加载前后的显存变化」或者「某个子模块的显存占用」。这时候就得自己写个监控器。

下面是我常用的一个轻量级监控类：

class MemoryMonitor:
    def __init__(self, device=0):
        self.device = device
        self.records = []

    def snapshot(self, tag=""):
        allocated = torch.cuda.memory_allocated(self.device)
        reserved = torch.cuda.memory_reserved(self.device)
        self.records.append({
            "tag": tag,
            "allocated_MB": allocated / 1024**2,
            "reserved_MB": reserved / 1024**2,
        })

    def report(self):
        print(f"{'Tag':<20} {'Allocated(MB)':>15} {'Reserved(MB)':>15}")
        print("-" * 50)
        for r in self.records:
            print(f"{r['tag']:<20} {r['allocated_MB']:>15.2f} {r['reserved_MB']:>15.2f}")

# 使用示例
monitor = MemoryMonitor()
monitor.snapshot("before_model")
model = MyModel().cuda()
monitor.snapshot("after_model")
optimizer = torch.optim.Adam(model.parameters())
monitor.snapshot("after_optimizer")
monitor.report()

这个工具虽然简单，但很实用。你想想看，在模型加载、数据预处理、前向传播、反向传播这些关键节点各打一个快照，显存的「大头」在哪就一目了然了。

  进阶思路： 如果你需要更细粒度的监控，可以结合 torch.cuda.memory_stats() 获取完整的统计字典，里面包含分配次数、释放次数、碎片率等几十个指标。我一般会在训练结束后，把 memory_stats 和 memory_summary 一起 dump 到文件里，方便事后分析。

知识体系总览

下面这张图总结了本章的核心逻辑：从「快速查看」到「深度分析」，再到「自定义扩展」，形成一个完整的监控链路。

嗯，以上就是显存监控的四个层次。从最基础的 nvidia-smi 到自定义工具，每个都有它的适用场景。我个人建议：日常训练用 torch.cuda 就够了，出问题时再上 Profiler。别一上来就全上，反而把自己搞晕了。

想要更多资料，联系可可

📚 课程目录

2. 显存监控：nvidia-smi、torch.cuda、PyTorch Profiler 与自定义工具

2.1 nvidia-smi：最基础的显存快照

2.2 torch.cuda：Python 里的实时监控

2.3 PyTorch Profiler：深度分析利器

2.4 自定义显存监控工具

知识体系总览