2、性能分析工具链:Linux perf、Nsight Systems、Nsight Compute、PyTorch Profiler、TensorBoard Profiling
做推理优化这些年,我最大的感触就是:没有趁手的工具,你就是在黑夜里摸瞎。很多同学一上来就盯着代码改,改完发现没效果,甚至更慢了。说白了,你得先知道瓶颈在哪,才能对症下药。
这一章,我就把平时最常用的几把「手术刀」给你捋一遍。它们各有各的脾气,也各有各的绝活。
核心观点:没有万能工具,只有最合适的组合。CPU瓶颈找perf,GPU瓶颈找Nsight,框架层面找PyTorch Profiler。
2.1 Linux perf —— CPU 瓶颈的照妖镜
先聊聊 Linux perf。这玩意儿是内核自带的,免费、轻量、无侵入。我个人习惯在排查 CPU 瓶颈时第一个用它。
它的核心能力就两个:采样和计数。采样告诉你程序时间花在哪了,计数告诉你硬件事件(比如缓存未命中、分支预测失败)发生了多少次。
我的经验:很多同学以为 perf 只能看热点函数。其实它最厉害的是看 cache-misses 和 branch-misses。我曾经遇到一个模型,算子本身不慢,但数据排布导致 L2 缓存疯狂 miss,性能直接腰斩。perf 一跑,问题就暴露了。
基本用法很简单:
# 采样热点
perf record -g -p <pid> -- sleep 10
perf report -g graph
# 统计硬件事件
perf stat -e cycles,instructions,cache-misses,branch-misses ./your_app
这里有个坑要注意:-g 参数会记录调用栈,但采样频率太高会影响程序行为。我一般设 99Hz 或 199Hz,够用又不至于太干扰。
避坑指南:我曾经在容器里用 perf,发现采不到内核态的事件。后来才意识到需要 --privileged 权限。另外,虚拟化环境(比如 AWS 的某些实例)可能不支持硬件 PMU 计数,这时候 perf stat 的结果就不准了。
2.2 Nsight Systems —— GPU 时间线的上帝视角
说到 GPU 分析,Nsight Systems 是我最常用的工具。它不像 perf 那样只给你一堆数字,而是给你一个可视化的时间线。你一眼就能看出 CPU 在等 GPU,还是 GPU 在等 CPU,还是数据传输在拖后腿。
我建议你重点关注这几个维度:
- CPU vs GPU 活动重叠度 —— 理想情况是 CPU 在提交下一个 batch 时,GPU 正在算上一个 batch
- CUDA 流之间的依赖 —— 不必要的同步点会杀死并行度
- 内存传输(H2D / D2H) —— 很多时候瓶颈不在计算,而在 PCIe 带宽
启动方式也很直接:
# 命令行模式
nsys profile -o output_file -t cuda,nvtx,osrt python your_script.py
# 或者用 GUI 模式
nsys-ui
嗯,这里要注意:Nsight Systems 的采样开销其实不小。我见过有人开着全量追踪跑大模型,结果 profiler 本身占了 30% 的 GPU 时间。所以生产环境上,我一般只开 -t cuda,nvtx,关掉 OS runtime 追踪。
2.3 Nsight Compute —— Kernel 级别的显微镜
如果说 Nsight Systems 是看全局的,那 Nsight Compute 就是看局部的。它专门用来分析单个 CUDA kernel 的性能。
你想想看,一个 kernel 跑得慢,到底是计算瓶颈还是访存瓶颈?是寄存器压力太大还是 shared memory 用得太少?Nsight Compute 会给你一个详细的性能分析报告,包括 occupancy、memory throughput、compute throughput 等等。
我常用的命令:
# 分析单个 kernel
ncu --kernel-name your_kernel -o report_file python your_script.py
# 或者交互式分析
ncu-ui report_file.ncu-rep
关键指标解读:
- Memory Throughput vs Compute Throughput —— 哪个接近理论峰值,哪个就是瓶颈
- Occupancy —— 不是越高越好,但太低肯定有问题
- Warp Stall Reasons —— 告诉你 warp 在等什么(等待数据?等待指令?)
我记得有一次调一个 attention kernel,Nsight Compute 显示 memory throughput 只有 30%。我一看,原来是数据排布导致合并访问没做好。改成 row-major 后,吞吐直接翻倍。
2.4 PyTorch Profiler —— 框架层面的瑞士军刀
如果你用的是 PyTorch,那 PyTorch Profiler 是你绕不开的工具。它跟框架深度集成,能告诉你每个算子的耗时、内存分配、甚至 CUDA kernel 的启动时间。
用法非常 Pythonic:
import torch.profiler as profiler
with profiler.profile(
activities=[profiler.ProfilerActivity.CPU, profiler.ProfilerActivity.CUDA],
record_shapes=True,
with_stack=True
) as prof:
output = model(input_tensor)
prof.step()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
我个人最常用的是 record_shapes=True。为什么?因为很多性能问题跟 tensor shape 有关。比如某个算子输入是 [1, 3, 224, 224] 时很快,但变成 [4, 3, 224, 224] 时突然变慢,这往往是因为 kernel 的 tile size 没对齐。
小技巧:用 prof.export_chrome_trace("trace.json") 导出 trace 文件,然后拖到 Chrome 的 chrome://tracing 里看。那个时间线比命令行输出直观多了。
2.5 TensorBoard Profiling —— 可视化对比的利器
最后聊聊 TensorBoard Profiling。它其实是 PyTorch Profiler 的「可视化升级版」。你可以在 TensorBoard 里看到训练/推理的完整时间线,还能做多轮对比。
我一般这么用:
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir="./logs")
with profiler.profile(...) as prof:
# 你的推理代码
pass
writer.add_profiler(prof)
writer.close()
# 然后启动 TensorBoard
# tensorboard --logdir ./logs
在 TensorBoard 里,我最喜欢的功能是 「Trace Viewer」 和 「Memory Profiling」。前者能让你像看视频一样逐帧分析算子执行顺序,后者能帮你发现内存泄漏或者碎片化问题。
避坑指南:我曾经在分布式训练场景下用 TensorBoard Profiling,发现不同 rank 的 trace 时间戳对不上。后来才意识到需要设置 torch.distributed.barrier() 来同步。另外,trace 文件太大时(超过几百 MB),TensorBoard 会卡死。建议只记录关键 step。
小结
这五个工具,各有各的定位。我个人的工作流是这样的:
- 先用 PyTorch Profiler 快速定位哪个算子最慢
- 如果是 CPU 算子,用 Linux perf 深挖热点函数和缓存行为
- 如果是 GPU 算子,用 Nsight Systems 看整体时间线,确认是计算、访存还是同步问题
- 再用 Nsight Compute 对具体 kernel 做微观分析
- 最后用 TensorBoard Profiling 做可视化对比,验证优化效果
工具是死的,思路是活的。你用得多了,自然就知道什么时候该掏出哪把刀。