PyTorch Profiler

一句话定位：框架层的性能显微镜——把一次前向/反向拆成算子级明细，回答”时间花在哪个算子、是 CPU 还是 GPU、显存怎么涨的”。

1. 能采集什么

• CPU / GPU 时间线：每个算子在 CPU 侧（Python/调度）与 GPU 侧（Kernel 执行）分别耗时多少，可导出 Chrome trace 可视化查看；
• 算子耗时排序：按 self/total time 对算子聚合排序，直接给出 Top-N 热点；
• 显存占用趋势：记录分配/释放曲线与峰值，配合 OOM 排查（3.2）；
• Kernel 与 Python 栈关联：with_stack=True 可把某个 GPU Kernel 回溯到具体的 Python 代码行——这是它相对 Nsight（5.2）的最大优势：框架语义清晰、能定位到自己写的那行代码。

2. 关键指标

• GPU Kernel 占比：GPU Kernel 累计时间 / 总墙钟时间。占比低意味着大量时间耗在等待、拷贝或 Python 开销上 → 先解决供给与调度（见 3.3）；
• 算子 Top-N 耗时：优化必须从最贵的算子入手（遵循 Amdahl 定律），避免优化了占比 1% 的算子却期待整体提速。

3. 使用要点

• 一定要跳过预热步（前几个 iteration 含 CUDA 上下文初始化、autotune、缓存分配器预热），用 schedule 设置 wait/warmup/active 阶段，只统计稳定态（同 5.4）；
• 采样本身有开销，不要在采样状态下评估绝对性能；
• 定位到热点算子后，再决定动作：算子融合、换实现（如启用 FlashAttention，3.7）、调 batch、或量化（4）。

4. 与大数据工具的对应

PyTorch Profiler 的”算子耗时排序” ≈ Spark UI 的”Stage/Task 耗时明细”——都是先按耗时排序找热点，再针对热点提假设（详见 5.3 映射表）。

参考：PyTorch 官方 Profiler 教程