性能分析工具：用PyTorch Profiler和Nsight Systems揪出性能瓶颈
说实话，算子优化这件事，最怕的不是代码写不好，而是不知道慢在哪。
我见过太多人，上来就对着代码一顿魔改，改完一跑，反而更慢了。为什么？因为没有用数据说话。今天我们就聊聊怎么用两个趁手的工具——PyTorch Profiler和Nsight Systems，把性能瓶颈一个个揪出来。
先搞清楚：我们到底在分析什么？
性能分析，说白了就是回答三个问题：
- 时间花在哪了？
- 显存用在哪了？
- 计算单元吃饱了没？
我个人习惯，拿到一个模型，先跑一遍Profiler，看看整体情况。就像去医院体检，先做个全身扫描，再针对性地查具体项目。
核心原则：不要猜，要测。你觉得慢的地方，往往不是真正的瓶颈。
PyTorch Profiler：你的第一把手术刀
PyTorch自带的Profiler，说实话，够用了。我刚开始用的时候也觉得它输出太乱，后来发现是自己没用好。
基本用法
import torch
from torch.profiler import profile, record_function, ProfilerActivity
# 定义一个简单的模型
model = torch.nn.Linear(1024, 1024).cuda()
input_tensor = torch.randn(128, 1024).cuda()
# 开始性能分析
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True,
with_stack=True
) as prof:
for _ in range(10):
with record_function("model_inference"):
output = model(input_tensor)
# 打印结果
print(prof.key_averages().table(
sort_by="cuda_time_total",
row_limit=10
))这段代码跑完，你会看到一张表格，里面列出了每个操作的耗时、调用次数、显存占用。嗯，这里要注意，record_shapes=True这个参数很重要，它能告诉你每个张量的形状，方便你判断是不是形状不匹配导致的性能问题。
怎么看这张表？
| 列名 | 含义 | 我的经验 |
|---|---|---|
| Name | 操作名称 | 看是不是你预期的算子 |
| CPU total | CPU端总耗时 | 如果远大于CUDA total，说明CPU在拖后腿 |
| CUDA total | GPU端总耗时 | 这个才是GPU真正干活的时间 |
| Self CUDA | 算子自身耗时（不含子操作） | 重点关注这个，排除法找瓶颈 |
| Mem Usage | 显存占用 | 看有没有显存泄漏或者峰值过高 |
我曾经遇到过一个案例，模型训练特别慢，一看Profiler，发现有个torch.where操作占了40%的时间。换成torch.masked_select之后，速度直接翻倍。你看，没有Profiler，你根本不知道问题出在这。
小技巧：用prof.export_chrome_trace("trace.json")导出trace文件，然后在Chrome浏览器地址栏输入chrome://tracing，拖进去看火焰图。视觉上更直观，能一眼看出哪些操作是串行的。
Nsight Systems：深入GPU内部
PyTorch Profiler能告诉你哪个算子慢，但如果你想问「为什么这个算子慢」，就得请出Nsight Systems了。
Nsight Systems是NVIDIA官方的性能分析工具，能看到GPU内核的执行情况、内存带宽利用率、流多处理器（SM）占用率等等。说白了，它能告诉你GPU到底在忙什么，还是在摸鱼。
命令行使用
# 采集性能数据
nsys profile -o my_trace -t cuda,nvtx,osrt python train.py
# 生成可视化报告
nsys stats my_trace.qdrep我个人习惯，先用-t cuda,nvtx只采集GPU和NVTX标记的数据，这样文件小，分析快。等找到可疑点了，再加osrt看系统调用。
怎么看Nsight Systems的报告？
打开生成的.qdrep文件，你会看到这样的界面：
- CUDA API追踪：看CPU调用了哪些CUDA API，有没有同步等待
- GPU内核执行：看每个kernel的启动时间、执行时间、占用率
- 内存操作：看H2D（主机到设备）和D2H（设备到主机）的传输
我记得有一次，一个卷积算子跑得特别慢。用Nsight Systems一看，发现它的SM占用率只有30%。说白了，就是GPU的算力没喂饱。原因是输入通道数太少，卷积核太小，导致每个SM上的线程块数量不够。后来我改了一下数据布局，把多个小卷积合并成一个大卷积，占用率直接拉到80%以上。
避坑指南：我曾经在Nsight Systems上看到一个kernel的耗时特别长，以为是它的问题。后来仔细一看，发现是因为前面的kernel没执行完，这个kernel在排队等。所以，看耗时的时候，一定要区分「执行时间」和「墙上时间」。执行时间才是kernel真正干活的时间，墙上时间包含了等待。
实战：用两个工具定位一个真实瓶颈
我们来看一个实际案例。假设你有一个Transformer模型，训练速度不理想。我们来一步步排查。
第一步：PyTorch Profiler全局扫描
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True) as prof:
for batch in dataloader:
output = model(batch)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
prof.step()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))跑完之后，发现aten::softmax占了15%的时间，aten::bmm占了20%。嗯，这两个是注意力机制的核心操作，慢是正常的。但等等，还有一个aten::copy_占了10%？这个不对劲。
第二步：Nsight Systems深入分析
用Nsight Systems采集一次训练迭代的数据，重点关注copy_操作。你会发现，这个copy_其实是从GPU拷贝数据到CPU，或者反过来。为什么会这样？
我遇到过类似的情况，原因是某个操作返回了CPU张量，然后又被送回了GPU。说白了，就是数据在CPU和GPU之间来回倒腾，浪费了大量带宽。
第三步：定位根因并修复
检查代码，发现是torch.where(condition, x, y)这个操作，如果condition是CPU张量，它会先把x和y拷到CPU，算完再拷回GPU。解决方案很简单：把condition也放到GPU上。
# 修复前
condition = torch.tensor([True, False, True]) # 默认在CPU
result = torch.where(condition, x, y) # 触发拷贝
# 修复后
condition = condition.cuda() # 确保在GPU上
result = torch.where(condition, x, y) # 没有拷贝改完之后，再跑一次Profiler，copy_的占比从10%降到了0.5%。训练速度提升了8%。
总结一下我的排查思路：
- 先用PyTorch Profiler找到耗时Top的操作
- 对可疑操作，用Nsight Systems看GPU内部行为
- 结合代码，判断是算法问题还是实现问题
- 修复后，再用Profiler验证效果
两个工具的分工
最后，我用一张图总结一下这两个工具的关系：
说白了，PyTorch Profiler是「侦察兵」，帮你快速找到可疑目标。Nsight Systems是「特种兵」，深入敌后搞清楚具体情况。两个配合着用，没有搞不定的性能问题。
我的建议：刚开始接触性能分析，先别急着上Nsight Systems。先用PyTorch Profiler跑一遍，把Top 5的耗时算子列出来。如果这些算子本身是合理的（比如矩阵乘法、卷积），再考虑用Nsight Systems看能不能优化。如果发现一些莫名其妙的操作（比如频繁的copy_、to），那大概率是代码层面的问题，直接改代码就行。
好了，工具就介绍到这里。记住，性能分析不是一次性的工作，而是贯穿整个优化过程的习惯。每次改完代码，都跑一遍Profiler看看效果，慢慢你就会有感觉了。