2、显存瓶颈识别：三把利器定位问题

显存瓶颈这东西，看不见摸不着。你模型跑得慢，到底是计算卡住了，还是显存堵住了？我刚开始做大规模训练时，经常在这上面栽跟头。后来慢慢摸出套路——用好三把工具：nvidia-smi、Nsight Systems、PyTorch Profiler。今天我就把这套方法掰开揉碎讲给你。

2.1 第一把刀：nvidia-smi 快速扫描

nvidia-smi 是每个搞 GPU 的人最先接触的命令。但说实话，大部分人只会看一眼显存占用率，这远远不够。

我个人习惯这样用：

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

# 只看关键字段
nvidia-smi --query-gpu=index,timestamp,memory.used,memory.total,utilization.gpu,utilization.memory --format=csv

这里有个坑——显存占用高不等于显存瓶颈。我在项目中遇到过好几次：显存用了 90%，但 GPU 利用率只有 20%。这说明什么？数据搬运太慢，GPU 在干等。真正的瓶颈在 PCIe 带宽或者数据预处理上。

2.2 第二把刀：Nsight Systems 时间线分析

nvidia-smi 只能看快照，看不到时间线上的细节。Nsight Systems 就不一样了，它能告诉你每一毫秒 GPU 在干什么。

我一般这样用：

# 命令行采集
nsys profile -o trace_output -t cuda,nvtx,osrt python train.py

# 然后打开 GUI 看时间线
nsys-ui trace_output.qdrep

打开时间线后，我主要看三个东西：

1. CUDA Memory 操作：有没有频繁的 alloc/free？如果有，说明显存碎片化严重
2. Memcpy 带宽：H2D 和 D2H 的传输速率是否接近理论峰值？
3. Kernel 执行间隙：两个 kernel 之间有没有大段空白？空白在等什么？

举个例子。我记得有一次跑一个推荐模型，Nsight Systems 时间线显示：每 2ms 就有一个 4MB 的 D2H 拷贝。为什么？因为 PyTorch 默认会在每个 step 结束后把 loss 拷回 CPU。改成 .item() 延迟获取后，训练速度直接快了 30%。

2.3 第三把刀：PyTorch Profiler 逐层分析

前两把刀看的是全局，PyTorch Profiler 能深入到每个算子。说白了，它能告诉你哪一层、哪个操作在吃显存。

import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    for step in range(10):
        output = model(inputs)
        loss = loss_fn(output, targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

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

输出结果里，我最关注这几列：

列名	含义	我的判断标准
Self CUDA Mem	算子自身分配的显存	超过 1GB 就要警惕
Total CUDA Mem	算子及其子操作的总显存	看占比，前 5 个算子通常吃掉 80% 显存
CUDA Mem Size	每次分配的块大小	如果频繁出现 4KB 小分配，说明有碎片问题
Call Stack	调用栈信息	定位到具体代码行

2.4 三把工具的配合套路

这三把工具不是孤立的。我一般按这个流程走：

1. nvidia-smi 初筛：看显存占用率和 GPU 利用率是否匹配。不匹配，先查数据加载。
2. Nsight Systems 定位阶段：看时间线，找 Memcpy 瓶颈还是 Kernel 瓶颈。
3. PyTorch Profiler 深挖算子：定位到具体哪一层、哪个操作在吃显存。

下面这张图总结了三者的关系：

2.5 避坑指南

这些年我用这三把工具踩过不少坑，挑几个典型的说说：

• nvidia-smi 显存数字不准？ 我曾经发现 nvidia-smi 显示显存占用 8GB，但 PyTorch 报 OOM。后来才知道，nvidia-smi 显示的是已分配显存，PyTorch 的缓存分配器会预占一部分。用 torch.cuda.memory_summary() 看更准。
• Nsight Systems 采集时间太长？ 默认配置会采集所有 CUDA API 调用，数据量巨大。我一般用 -c cudaProfilerApi --capture-range cudaProfilerApi 只采集关键区间。
• PyTorch Profiler 影响训练速度？ 会的，尤其是 profile_memory=True 时。我建议只在调试时开，或者每 100 个 step 采样一次。