TorchInductor GPU 性能分析

本节列出了有用的命令和工作流程，可以帮助您深入了解 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 为例。

1. 运行模型的基准测试脚本

   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 对此至关重要。

2. 在输出日志中，查找以下行

   **Compiled module path:
   /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**
   

每个已编译模块都有一行。如果没有额外的图形断点，我们将在日志中看到两行，分别对应于前向图和后向图。

对于我们的示例命令，我们分别为前向图和后向图获得了以下已编译模块

• https://gist.github.com/shunting314/c2a4d8a28b00fcb5586d0e9d9bf77f9f

• https://gist.github.com/shunting314/48efc83b12ec3ead950052e4a0220b10

3. 现在我们可以深入了解每个已编译模块的性能。为了便于说明，我们选择前向图的模块。为了方便起见，我将其命名为 fwd.py。使用 -p 参数直接运行它

   **> python fwd.py -p**
   

请参阅此 示例 gist 中的完整输出日志。

在输出中，您可以注意到以下内容

• 我们为配置文件编写了一个 Chrome 跟踪文件，因此我们可以加载跟踪文件并与它进行交互。在日志中，查找以下行以找到跟踪文件的路径。

配置文件的 Chrome 跟踪写入 /tmp/compiled_module_profile.json

将跟踪文件加载到 Chrome 中（访问 chrome://tracing，并在 Chrome 浏览器中加载文件，如 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%

  • 其余为 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 中查找内核名称，并找到类似以下的注释

# 内核路径：/tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py

为了方便起见，我将其重命名为 k.py。这是此 文件 的粘贴。

k.py 是一个独立的 Python 模块，包含内核代码及其基准测试。

直接运行 k.py 会报告其执行时间和带宽

我们可以通过运行以下命令来检查 max-autotune 是否对该内核有帮助

**TORCHINDUCTOR_MAX_AUTOTUNE=1 python /tmp/k.py**

我们也可以暂时添加更多缩减启发式方法并再次运行脚本，以检查这是否有助于内核。