PyTorch Profiler:定位框架与算子开销¶
torch.profiler 位于 Python 代码与 CUDA kernel 之间,能够把 PyTorch 算子、CPU 活动和 CUDA 活动关联起来。它适合回答“模型的哪一部分最耗时”,也是从系统级现象继续下钻的重要一步。
它能回答什么¶
- 哪个 PyTorch 算子占用时间最多?
- 算子被调用了多少次,单次耗时是否稳定?
- 时间主要消耗在 Python、CPU 还是 GPU?
- CPU 算子对应启动了哪些 CUDA kernel?
torch.compile是否减少了算子数量或形成融合区域?matmul、attention、all_reduce、copy_等操作各占多少比例?
如果把 nsys 看作交通路况图,那么 torch.profiler 更像是在统计“哪一种车辆造成了拥堵”。
最小采集示例¶
import torch
from torch.profiler import ProfilerActivity, profile, record_function
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True,
with_stack=True,
) as prof:
with record_function("model_inference"):
run_inference()
print(prof.key_averages().table(
sort_by="self_cuda_time_total",
row_limit=20,
))
真实服务不宜无限制采集。可以使用 schedule 只记录少量稳定迭代,并导出 Chrome trace 或 TensorBoard trace 进行查看。
读结果时看哪些列¶
| 指标 | 含义 | 适合判断 |
|---|---|---|
| CPU total | 算子及其子调用的 CPU 总时间 | 整段调用的 CPU 开销 |
| CPU self | 不包含子调用的 CPU 时间 | 算子自身是否有明显 CPU 开销 |
| CUDA total | 算子关联的 CUDA 总时间 | 哪类算子占用 GPU 最多 |
| CUDA self | 不包含子调用的 CUDA 时间 | 当前算子自身的 GPU 开销 |
| Calls | 调用次数 | 是否存在大量细碎调用 |
总时间高不一定意味着单次执行慢:一个很短的算子如果被调用成千上万次,也可能成为主要开销。因此需要同时查看总耗时、平均耗时和调用次数。
CUDA Graph 下的限制¶
在 CUDA Graph 场景中,profiler 可能主要看到 graph replay,而无法像 eager 模式一样完整呈现 graph 内每个框架算子与 kernel 的边界。vLLM 开启 CUDA Graph 后,不应只依赖 torch.profiler 得出结论。
可采用以下方式交叉验证:
- 用 Nsight Systems 观察 replay 前后和 GPU 时间线;
- 在可控环境中暂时关闭 CUDA Graph,比较算子构成;
- 对确认有问题的 kernel 使用 Nsight Compute。
常见误区¶
- 在 warm-up 阶段统计。 初始化、编译和缓存建立会污染稳态结果。
- 只按总耗时排序。 还应关注 self time、调用次数和输入 shape。
- 采集窗口过长。 profiler 自身有开销,长时间采集会改变服务行为。
- 把相关性当成根因。 高耗时算子可能只是输入规模或上游调度问题的结果。