显存监控工具：nvidia-smi详解、torch.cuda内存API、显存碎片可视化

做分布式训练，显存就是命根子。我见过太多人一上来就调batch size，结果OOM（内存溢出），连训练都跑不起来，更别提优化了。原因很简单——你连显存怎么没的都不知道。所以这一章，咱们把监控工具彻底搞透。

我会从三个层面展开：硬件级的nvidia-smi、框架级的torch.cuda API、以及内存布局的碎片可视化。三个视角合起来，你就能准确回答三个问题：显存谁占的？占在哪？是不是碎成渣了？

3.1 nvidia-smi 详解

nvidia-smi 是 N 卡必备命令。说实话，大多数人只用 nvidia-smi 看一眼显存占用，然后抱怨不够用。这其实浪费了。它能给你颗粒度更高的信息。

最常用的命令

# 监控模式，1秒刷新
watch -n 1 nvidia-smi

# 显示全部 GPU 的详细指标
nvidia-smi -q

# 只显示显存和利用率，干净些
nvidia-smi --query-gpu=index,name,memory.total,memory.used,memory.free,utilization.gpu,utilization.memory --format=csv

# 查看指定 GPU 上的进程
nvidia-smi pmon -c 1

我习惯用 --query-gpu 加 CSV 输出，方便扔进脚本做日志。比如跑长训练时，每小时抓一次存起来，回头定位 OOM 时间点特别好用。

输出关键字段解释

字段	含义	说明
Memory-Usage	显存总量 / 已用 / 空闲	单位 MiB，注意是 MiB 不是 MB
GPU-Util	计算单元利用率（%）	低于 30% 可能瓶颈在数据加载
Volatile GPU-Util	瞬时利用率	老版本用，新版本已合并到 GPU-Util
Processes	占用 GPU 的进程列表	显示 PID、进程名、显存占用
FB Memory Usage	显存分块使用情况	只有 -q 模式下才出现

避坑指南

• 不要只看 GPU-Util 高就以为没瓶颈。 显存带宽饱和也能导致利用率低，你想想看。
• 注意 memory.used 和 memory.reserved 的区别。 nvidia-smi 报告的是实际物理分配，而 PyTorch 内部还有缓存池。
• 进程列表有时候不准。 我遇到过 CUDA 上下文被僵尸进程占着，但列表显示空的，这时候 fuser -v /dev/nvidia* 更靠谱。

3.2 torch.cuda 内存 API

如果说 nvidia-smi 是外部视角，那 torch.cuda 就是内窥镜。它能告诉你进程内部的显存是怎么分的：哪些是张量占的，哪些是缓存池，哪些被 PyTorch 预分配了。

核心函数

import torch

# 当前已分配字节（不含缓存池）
print(torch.cuda.memory_allocated())

# 当前缓存池总大小（包括分配+预留）
print(torch.cuda.memory_reserved())

# 最大峰值（可以设 reset_peak 重置）
print(torch.cuda.max_memory_allocated())

# 漂亮的摘要，打印到stderr
torch.cuda.memory_summary()

# 获取当前设备
print(torch.cuda.current_device())

# 清空缓存（慎用）
torch.cuda.empty_cache()

分布式场景下的用法

多卡训练时，记得在每个 rank 里打印自己的显存。我一般这样写：

if torch.distributed.get_rank() == 0:
    print(f"Rank {rank}: allocated={torch.cuda.memory_allocated()/1e9:.2f}GB, "
          f"reserved={torch.cuda.memory_reserved()/1e9:.2f}GB")

这样输出不会乱，而且能对比不同卡之间的显存差异。如果某张卡明显偏高，八成是数据不平衡或者模型并行切分有问题。

缓存池是怎么回事？

PyTorch 默认使用 CUDA 缓存分配器，一次性从 GPU 申请一大块 memory_reserved，然后内部再切小块给你用。这样做的好处是减少 cudaMalloc 系统调用，坏处是显存碎片。你想想看，模型加载后释放一些中间变量，缓存池里就会出现空洞，后续申请大张量时就有可能 OOM，即使总空闲显存足够。

3.3 显存碎片可视化

说到碎片，光看数字不够，最好能把内存布局画出来。PyTorch 1.8+ 提供了 torch.cuda.memory_snapshot()，可以获取堆上每个分配的详细记录。配合 torch.cuda.memory._dump_snapshot 还能生成可视化。我这几年排查 OOM 问题，十有六七是碎片导致的。

获取快照

import torch.cuda

# 记录快照（仅在PyTorch DEBUG模式下可用）
snapshot = torch.cuda.memory_snapshot()

# 也可以保存为文件，供可视化工具使用
with open("/tmp/mem_snapshot.pkl", "wb") as f:
    import pickle
    pickle.dump(snapshot, f)

这个快照是个列表，每个元素包含 address、size、stream、history 等字段。你可以遍历它，统计最大的空闲空洞有多大。

实用分析方法

def analyze_fragmentation(snapshot):
    # 按地址排序
    sorted_blocks = sorted(snapshot, key=lambda x: x['address'])
    total_reserved = torch.cuda.memory_reserved()
    used = 0
    largest_gap = 0
    total_gaps = 0
    prev_end = 0
    
    for block in sorted_blocks:
        addr = block['address']
        size = block['size']
        if addr > prev_end:
            gap = addr - prev_end
            total_gaps += gap
            if gap > largest_gap:
                largest_gap = gap
        used += size
        prev_end = addr + size
    
    # 末尾到reserved_end的间隙
    reserved_end = max(b['address'] + b['size'] for b in sorted_blocks) if sorted_blocks else 0
    final_gap = total_reserved - reserved_end if total_reserved > reserved_end else 0
    total_gaps += final_gap
    largest_gap = max(largest_gap, final_gap)
    
    frag_ratio = total_gaps / total_reserved if total_reserved > 0 else 0
    return {
        'used': used,
        'free_in_reserved': total_reserved - used,
        'total_gaps': total_gaps,
        'largest_gap': largest_gap,
        'frag_ratio': frag_ratio
    }

# 用的时候：
snap = torch.cuda.memory_snapshot()
stats = analyze_fragmentation(snap)
print(f"碎片率: {stats['frag_ratio']*100:.1f}%")

环境变量调优碎片

# 限制保留内存比例（默认0.5）
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# 或设置最大分隔段数量
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

我建议在训练脚本开头加上：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True, max_split_size_mb:128'

这个配置能显著缓解碎片问题。但要注意，expandable_segments 需要 CUDA 11.8+ 和 PyTorch 2.0+。如果版本不够，碎片率会高不少。

推荐碎片可视化工具

• PyTorch Memory Visualizer：官方提供，用 torch.cuda.memory._dump_snapshot 生成 JSON，拖进 Chrome 看火焰图。
• nvidia-smi 的 FB Memory：nvidia-smi -q -d MEMORY 能看到桶状分配图，虽然粗糙但有参考价值。
• 自定义脚本：就像上面那个，用 memory_snapshot() 自己分析，适合自动化监控。

本章小结