TorchInductor GPU 性能分析 ¶

本节列出了有用的命令和工作流程，可以帮助您深入了解 TorchInductor 中模型的性能。当模型运行速度没有达到预期时，您可能想检查模型的各个内核。通常，占用大部分 GPU 时间的内核是最有趣的。之后，您还可以直接运行单个内核并检查其性能。PyTorch 提供了覆盖上述所有内容的工具。

相关环境变量 ¶

您可以在分析中使用以下环境变量：

• TORCHINDUCTOR_UNIQUE_KERNEL_NAMES

  • 默认情况下，TorchInductor 将 Triton 内核命名为 ‘triton\_’ 。当启用此环境变量时，inductor 会在跟踪中生成更有意义的内核名称，例如 triton_poi_fused_cat_155 ，其中包含内核类别（ poi 表示 pointwise）和原始 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 浏览器中访问 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 ，这是最昂贵的 reduction 内核，占整个前向图总墙时的 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**

我们也可能暂时添加更多的缩减启发式算法，然后再次运行脚本以检查这对内核的帮助。