TorchInductor GPU 性能分析#
创建于: 2023年7月28日 | 最后更新于: 2025年6月10日
本节列出了有助于深入了解 TorchInductor 中模型性能的有用命令和工作流。当模型运行速度不如预期时,您可能需要检查模型的各个内核。通常,占用大部分 GPU 时间的内核是最值得关注的。之后,您可能还想直接运行各个内核并检查其性能。PyTorch 提供了涵盖上述所有内容的工具。
相关的环境变量#
您可以在分析中使用以下环境变量
TORCHINDUCTOR_UNIQUE_KERNEL_NAMES
默认情况下,TorchInductor 将 Triton 内核命名为
‘triton\_’
。当启用此环境变量时,Inductor 会在跟踪中生成更有意义的内核名称,例如triton_poi_fused_cat_155
,其中包含内核类别(poi
代表逐点操作)和原始 ATen 运算符。此配置默认禁用,以提高编译缓存命中率。
TORCHINDUCTOR_BENCHMARK_KERNEL
启用此选项将使 Inductor 代码生成器能够对各个 Triton 内核进行基准测试。
TORCHINDUCTOR_MAX_AUTOTUNE
Inductor 自动调优器将对更多
triton.Configs
进行基准测试,并选择性能最佳的配置。这将增加编译时间,但有望提高性能。
分解模型 GPU 时间#
以下是将模型执行时间分解到各个内核的步骤。我们以 mixnet_l
为例。
为模型运行基准测试脚本
TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 TORCHINDUCTOR_BENCHMARK_KERNEL=1 python -u benchmarks/dynamo/timm_models.py –backend inductor –amp –performance –dashboard –only mixnet_l –disable-cudagraphs –training
注意
该工具依赖于内核名称来决定其类别。启用
TORCHINDUCTOR_UNIQUE_KERNEL_NAMES
对此至关重要。在输出日志中,查找以下行
**Compiled module path: /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**
我们为每个编译的模块有一行。如果没有额外的图中断,我们将在日志中看到 2 行这样的内容,一行用于前向图,一行用于后向图。
对于我们的示例命令,我们分别获得前向图和后向图的编译模块如下
现在我们可以深入了解每个独立编译模块的性能。为了说明方便,我们选择前向图的模块。我将其命名为
fwd.py
。使用-p
参数直接运行它**> python fwd.py -p**
在此 示例 gist 中查看完整的输出日志
在输出中,您会注意到以下内容
我们为性能分析写入了一个 chrome trace 文件,因此我们可以加载该跟踪并与之交互。在日志中,查找如下的行以找到 trace 文件的路径。
性能分析的 Chrome trace 已写入 /tmp/compiled_module_profile.json
将跟踪加载到 Chrome 中(访问 chrome 浏览器中的 chrome://tracing 并按照 UI 的提示加载文件)将显示如下 UI
您可以放大和缩小来检查性能分析。
我们通过类似以下的日志行报告 GPU 忙碌时间占总时间的百分比
GPU 忙碌时间百分比:102.88%
有时您可能会看到大于 100% 的值。原因在于 PyTorch 在启用性能分析时使用内核执行时间,而在禁用性能分析时使用总时间。性能分析可能会稍微扭曲内核执行时间。但总体来说应该不是大问题。
如果我们像这样运行
densenet121
模型,并使用较小的批量大小,我们会看到 GPU 忙碌时间百分比较低(Forward graph) Percent of time when GPU is busy: 32.69%
这表明模型存在大量 CPU 开销。这与启用 cudagraphs 大大提高了 densenet121 的性能这一事实一致。
我们可以将 GPU 时间分解到不同类别的内核。在
mixnet_l
示例中,我们看到逐点内核占 28.58%
归约内核占 13.85%
持久化归约内核占 3.89%
其余为用于 mm/conv 的 cutlass/cudnn 内核,占 56.57%
此信息可以在每个内核类别的报告的摘要行(最后一行)中找到。
我们还可以深入到某一类内核。例如,让我们检查归约内核
我们可以看到每个独立归约内核执行时间的排序表。我们还可以看到内核执行的次数。这在几个方面很有帮助
如果一个内核只占用很小的时间,例如 0.1%,改进它最多只能带来 0.1% 的总体收益。不值得花费大量精力。
如果一个内核占 2% 的时间,将其改进 2 倍可以带来 1% 的总体收益,这证明了付出的努力是值得的。
基准测试单个 Triton 内核#
假设我们想仔细查看 triton_red_fused\__native_batch_norm_legit_functional_16
,它是最昂贵的归约内核,在前向图的总体运行时间中占 2.19%。
我们可以在 fwd.py
中查找内核名称,并找到类似这样的注释
# kernel path: /tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py

为了方便起见,我将其重命名为 k.py。这是此 文件 的粘贴内容。
k.py
是一个独立的 Python 模块,包含内核代码及其基准测试。
直接运行 k.py
将报告其执行时间和带宽

我们可以通过运行以下命令来检查 max-autotune 是否对此内核有帮助
**TORCHINDUCTOR_MAX_AUTOTUNE=1 python /tmp/k.py**
我们还可以暂时添加更多归约启发式规则并重新运行脚本,以检查这些规则如何帮助内核。