把工具搞混最常见的错误是:用 ncu 去找 DataLoader 瓶颈, 用 py-spy 去看 kernel 是不是 memory-bound。 四个工具的关系是不同抽象层次的显微镜—— 先在高层用粗粒度工具定位区域,再用低层细粒度工具放大问题。

nsys 把整个进程的 CPU 线程、CUDA stream、kernel 启动、memcpy、NCCL 调用 画在同一条时间线上,几乎是判断"GPU 究竟在不在干活"的唯一办法。 它不告诉你某个 kernel 内部好不好,只告诉你整个系统的协作是否顺畅。

常用命令

# 1) 最常用:profile 一次完整的 Python 训练脚本
nsys profile -o resnet \
     --trace=cuda,nvtx,cudnn,cublas \
     --cuda-memory-usage=true \
     python train.py

# 2) 只 profile 中间 5 秒(避开 warmup),减少文件大小
nsys profile -o resnet -y 10 -d 5 python train.py

# 3) 多 GPU / 多进程:每个 rank 输出独立 .nsys-rep
nsys profile -o rank%q{LOCAL_RANK} \
     --trace=cuda,nvtx,nccl \
     torchrun --nproc_per_node=4 train.py

# 4) 命令行查看摘要(无 GUI 时)
nsys stats resnet.nsys-rep
# 输出:CUDA Kernel Time、Memcpy Stats、NVTX ranges、CPU sampling……

# 5) GUI 打开
nsys-ui resnet.nsys-rep

用 NVTX 给 timeline 加标签

默认 timeline 只有 CUDA API 和 kernel 名字,信息有限。 用 torch.cuda.nvtx.range_push/pop 给关键代码段打标签, timeline 上会出现彩色长条,极大提高可读性。

import torch
from torch.cuda import nvtx

for step, batch in enumerate(loader):
    nvtx.range_push(f"step_{step}")

    nvtx.range_push("forward")
    out = model(batch)
    loss = criterion(out, target)
    nvtx.range_pop()

    nvtx.range_push("backward")
    loss.backward()
    nvtx.range_pop()

    nvtx.range_push("optimizer")
    optimizer.step()
    optimizer.zero_grad()
    nvtx.range_pop()

    nvtx.range_pop()  # step_{step}

Day 05 已经用过 ncu。今天补充三个进阶用法: 指定 kernel 抓取(避免 profile 整个程序)、 Source View(把指标对齐到 CUDA 源码行)、 两份报告对比(优化前后的定量评估)。

最该看的几个 Section

优化前后对比工作流

# Step 1: 编译时保留行号信息
nvcc -O2 -lineinfo -o tiled_gemm tiled_gemm.cu

# Step 2: profile baseline 与优化版各一份
ncu --set full --import-source on -k tiled_gemm \
     -o baseline ./tiled_gemm

# 改完代码再跑一次
ncu --set full --import-source on -k tiled_gemm \
     -o optimized ./tiled_gemm

# Step 3: 在 ncu-ui 中加载 optimized.ncu-rep,
#         然后 File → Add Baseline → 选 baseline.ncu-rep
#         所有指标会显示 ± 百分比变化
ncu-ui optimized.ncu-rep

nsys 看不清 Python 里发生什么——它能告诉你"CPU 在某段时间闲着", 但不知道是哪段 Python 代码在闲。py-spy 填补这个空隙: 不修改一行代码、对运行中的 Python 进程做采样, 生成火焰图。AI Infra 场景下最常用来定位 DataLoader 瓶颈。

三个最有用的命令

# 1) 对运行中的训练进程实时 top(无侵入)
py-spy top --pid 12345

# 2) 采样 60 秒,生成火焰图
py-spy record -o flame.svg --pid 12345 --duration 60

# 3) 训练挂死时,看每个线程在哪
py-spy dump --pid 12345

# 4) 直接启动并 profile(适合短脚本)
py-spy record -o flame.svg -- python train.py

# 5) DataLoader worker 也要 profile,加 --subprocesses
py-spy record -o flame.svg --subprocesses -- python train.py

torch.profiler 是 PyTorch 内置的 profiler,介于 nsys 与 ncu 之间—— 它同时记录 op 调用(forward/backward/optimizer) 和 底层 CUDA kernel。 最大的优势:数据用 TensorBoard 打开,Operator View / Kernel View / Memory View 一站式查看。

最小可用模板

import torch
from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity

prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(
        wait=1,        # 跳过前 1 个 step (warmup)
        warmup=1,      # 预热 1 个 step (不记录但激活 profiler)
        active=3,      # 真正采集 3 个 step
        repeat=1,      # 整个周期重复 1 次
    ),
    on_trace_ready=tensorboard_trace_handler("./log/resnet"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
)

with prof:
    for step, batch in enumerate(loader):
        out = model(batch)
        loss = criterion(out, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        prof.step()   # 重要!推进 profiler 的 schedule 阶段
        if step >= 10:
            break

# 命令行查看 top kernels
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

用 TensorBoard 查看

# 安装 TensorBoard Plugin
pip install torch-tb-profiler

# 启动 TensorBoard 指向 trace 目录
tensorboard --logdir=./log

# 浏览器打开后切到 "PYTORCH_PROFILER" 标签页

把 nsys / ncu / py-spy / torch.profiler 串起来

四个工具不是互斥的,而是同一个调查里的不同放大倍数。 一个完整的性能问题诊断通常这样走:

Lab — profile_resnet.py

用 ResNet-50 + CIFAR-10(或 ImageNet 子集)写一个故意有问题的训练脚本—— num_workers=0、每 step 打印 loss.item()、 用 PIL 现场 resize——然后用今天的四个工具找出每个问题。

# profile_resnet.py — 故意写慢的训练脚本
import torch, torchvision
from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
from torch.cuda import nvtx

model = torchvision.models.resnet50().cuda()
opt   = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = torch.nn.CrossEntropyLoss()

# 故意:num_workers=0(单进程加载,严重瓶颈)
loader = torch.utils.data.DataLoader(dataset, batch_size=64,
                                      num_workers=0, pin_memory=False)

prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3),
    on_trace_ready=tensorboard_trace_handler("./log/resnet"),
    record_shapes=True, profile_memory=True, with_stack=True,
)

with prof:
    for step, (x, y) in enumerate(loader):
        nvtx.range_push(f"step_{step}")
        x, y = x.cuda(), y.cuda()

        nvtx.range_push("forward");   out  = model(x);            nvtx.range_pop()
        nvtx.range_push("loss");      loss = loss_fn(out, y);     nvtx.range_pop()
        nvtx.range_push("backward");  loss.backward();            nvtx.range_pop()
        nvtx.range_push("opt");       opt.step(); opt.zero_grad(); nvtx.range_pop()

        # 故意:每步打印,触发隐式 D2H 同步
        print(f"step {step}: loss={loss.item():.4f}")

        prof.step()
        nvtx.range_pop()
        if step >= 10: break

四步体检流程

# Step 1: torch.profiler — 看 op 级别占比
python profile_resnet.py
tensorboard --logdir=./log
# 在 Operator View 看 GPU 利用率,大概率 < 50%

# Step 2: nsys — 看时间线上 GPU 空隙
nsys profile -o resnet --trace=cuda,nvtx -y 5 -d 10 \
     python profile_resnet.py
nsys stats resnet.nsys-rep
# 应该能看到大块 GPU idle,以及 print/.item() 引起的 sync

# Step 3: py-spy — 看 CPU 端 Python 热点
py-spy record -o flame.svg --subprocesses --duration 30 \
       -- python profile_resnet.py
# 在火焰图里看 PIL / collate / to_tensor 占多少

# Step 4: ncu — 钻最热的 kernel(可选)
ncu --set full --launch-skip 50 --launch-count 5 \
     -k regex:ampere_sgemm -o resnet_gemm \
     python profile_resnet.py

预期诊断结果

1. 画一张图:把 nsys、ncu、py-spy、torch.profiler 按"看哪一层"排列,并写出各自的适用场景。
2. 训练脚本 GPU 利用率只有 40%,你的诊断顺序是什么?第一个动用的工具是什么?第一个排除的可能性是什么?
3. 解释为什么 loss.item() 出现在训练循环里会让 nsys timeline 出现规律性 GPU idle。
4. 给你的 Day 05 tiled_gemm 跑一次 ncu,记录 SOL 卡片的 Compute % 和 Memory %,说出它是 compute-bound 还是 memory-bound,以及下一步应该往哪个方向优化。
5. 写一个三行的"profile 决策表":出现 X 现象 → 用 Y 工具 → 看 Z 指标。覆盖至少 5 种现象。