分析以了解 torch.compile 的性能¶

使用 torch.profiler 的目的¶

torch.profiler 有助于以内核级粒度了解程序的性能 - 例如，它可以显示程序级图中断和 GPU 利用率。分析器提供的数据通常可以帮助用户了解在哪里进一步调查以了解模型性能。

要了解内核级性能，还存在其他工具。可以使用 NVIDIA 的 ncu 工具，或inductor 的分析工具。

使用 torch.profiler 和查看跟踪的基本知识¶

示例程序：我们将使用此 resnet18 分析示例。请注意此示例程序的以下部分

包含预热运行以等待编译完成（这将预热 CUDA 缓存分配器等系统）
使用torch.profiler.profile() 上下文来分析我们感兴趣的部分
使用prof.export_chrome_trace("trace.json") 导出分析工件。

import torch
from torchvision.models import resnet18

model = resnet18().cuda()
inputs = [torch.randn((5, 3, 224, 224), device='cuda') for _ in range(10)]

model_c = torch.compile(model)

def fwd_bwd(inp):
    out = model_c(inp)
    out.sum().backward()

# warm up
fwd_bwd(inputs[0])

with torch.profiler.profile() as prof:
    for i in range(1, 4):
        fwd_bwd(inputs[i])
        prof.step()

prof.export_chrome_trace("trace.json")

查看 Chrome 跟踪：在 Chrome 浏览器中，打开 chrome://tracing 并加载 json 文件。使用“w”和“s”键放大和缩小，使用“a”和“d”键向左和向右滚动。“？” 将显示一个“帮助”屏幕，其中包含快捷键列表。

Example of a basic chrome trace, visualized in the chrome://tracing viewer

在这里，我们观察到：* CompiledFunction 和 CompiledFunctionBackward 事件，它们对应于 dynamo 编译的区域。* CPU 事件位于顶部，GPU 事件位于底部。

CPU 和 GPU 事件之间的流程

GPU 上的每个内核都是在 CPU 上运行的代码启动后发生的。分析器可以在 GPU 和 CPU 事件之间绘制连接（即“流程”）以显示哪个 CPU 事件启动了 GPU 内核。这特别有用，因为除了少数例外情况外，GPU 内核是异步启动的。

要查看流程连接，请单击 GPU 内核并单击“ac2g”

Visualization in the chrome://trace viewer, showing an async flow between a kernel and its launching location.

或者，使用顶部的“流程事件”下拉菜单打开所有流程。

解决 CUDA 图形分析问题¶

启用 CUDA 图形后，某些 cuda 配置（驱动程序版本低于 525.85.12 或 CUDA < 12）可能会在分析工具和 CUDA 图形之间遇到问题。要解决这些问题，请在程序顶部添加一个空的分析上下文

import torch

torch.profiler._utils._init_for_cuda_graphs()

# ... rest of program

了解编译时间¶

要了解为什么编译需要很长时间，您可以分析 torch.compile 程序的第一次调用。请记住，编译的分析跟踪可能比典型分析更失真，因为编译工作负载可能与典型的 PyTorch 工作负载大不相同。在某些情况下，跟踪文件也可能非常大。跟踪文件 > 1GB 可能难以使用 Chrome 跟踪工具打开。

注意：大致相同的信息也可以使用torch._dynamo.utils.compile_times() 以非图形格式获取。此实用程序不会显示编译步骤发生的时间，但它会显示每个步骤花费的时间 - 并且时间不会受到任何分析开销的影响。

请参见下面的示例

import torch
from torchvision.models import resnet18

model = resnet18().cuda()
inputs = [torch.randn((5, 3, 224, 224), device='cuda') for _ in range(10)]

model_c = torch.compile(model)

def fwd_bwd(inp):
    out = model_c(inp)
    out.sum().backward()

def warmup_compile():
    def fn(x):
        return x.sin().relu()

    x = torch.rand((2, 2), device='cuda', requires_grad=True)
    fn_c = torch.compile(fn)
    out = fn_c(x)
    out.sum().backward()

with torch.profiler.profile() as prof:
    with torch.profiler.record_function("warmup compile"):
        warmup_compile()

    with torch.profiler.record_function("resnet18 compile"):
        fwd_bwd(inputs[0])

prof.export_chrome_trace("trace_compile.json")

A visualization in the chrome://trace viewer, showing dynamo and inductor compilation steps

注意以下几点

第一次调用应在分析期间发生，以便捕获编译
添加预热编译以初始化任何需要延迟初始化的系统。

查找图中断：“Torch-Compiled Region”和“CompiledFunction”¶

虽然存在用于识别图中断的日志记录工具，但分析器提供了一种快速识别图中断的可视化方法。有两个分析器事件需要查找：**Torch-Compiled Region** 和 **CompiledFunction**。

**Torch-Compiled Region** - 在 PyTorch 2.2 中引入 - 是一个覆盖整个编译区域的分析器事件。图中断几乎总是看起来相同：嵌套的“Torch-Compiled Region”事件。

如果您分别对两个单独的函数应用 torch.compile()，则通常应该期望看到两个相邻的（即未堆叠/嵌套）Torch-Compiled 区域。同时，如果您遇到图中断（或已禁用/跳过的区域），则应期望嵌套的“Torch-Compiled Region”事件。

**CompiledFunction** - 在 PyTorch 2.0 中引入 - 是一个分析器事件，当任何输入需要梯度时出现。每个图中断都会中断 CompiledFunction 块，将其分成两部分。CompiledFunction 事件仅在涉及 Autograd 时出现，即图的一些输入张量的 requires_grad=True。

当 CompiledFunction 出现在跟踪中时，它通常会在反向传递中与 CompiledFunctionBackward 事件配对。如果调用了反向函数，则跟踪中应该会出现连接这两者的“前向-反向链接”。

如果您的用例包含一个不需要梯度且不包含“Torch-Compiled Region”事件的图，则可能难以确定是否正确应用了 torch.compile。一个线索可能是 Inductor 生成的 Triton 内核的存在。

请参见下面的合成示例以进行演示

import torch
import torch._dynamo

class ModelWithBreaks(torch.nn.Module):
    def __init__(self):
        super().__init__()
        def create_sequential():
            return torch.nn.Sequential(
                torch.nn.Linear(128, 128),
                torch.nn.ReLU(),
                torch.nn.Linear(128, 128),
                torch.nn.ReLU(),
            )
        self.mod1 = create_sequential()
        self.mod2 = create_sequential()
        self.mod3 = create_sequential()
        self.mod4 = create_sequential()

    def forward(self, inp):
        mod1 = self.mod1(inp)
        torch._dynamo.graph_break()
        mod2 = self.mod2(mod1)
        torch._dynamo.graph_break()
        mod3 = self.mod3(mod2)
        torch._dynamo.graph_break()
        mod4 = self.mod4(mod3)
        return mod4


model = ModelWithBreaks().cuda()
inputs = [torch.randn((128, 128), device='cuda') for _ in range(10)]

model_c = torch.compile(model)

def fwd_bwd(inp):
    out = model_c(inp)
    out.sum().backward()

# warm up
fwd_bwd(inputs[0])

with torch.profiler.profile() as prof:
    for i in range(1, 4):
        fwd_bwd(inputs[i])
        prof.step()

prof.export_chrome_trace("trace_break.json")

Visualization in the chrome://trace viewer, showing nested Torch-Compiled Region events and multiple CompiledFunction events - indicating graph breaks.

运算符内核¶

当启动运算符时，我们期望看到一些事件

CPU 端事件
内核启动（如果处理 GPU 内核）
GPU 端事件

Visualization in the chrome://trace viewer, showing the three types of events: CPU-side event, kernel launch, and GPU-side event

由 Inductor 生成的 Triton 内核：1. CPU 端事件应显示为以“triton_”为前缀的事件。当前这些事件的信息量很少 - 内核名称和启动，但比典型的 aten 内核启动（包含输入形状、类型等）信息量少。2. 内核启动应显示为 cuLaunchKernel 而不是 cudaLaunchKernel（cudaLaunchKernel 通常用于 aten 操作）3. GPU 端事件应出现，其名称的描述性将取决于 inductor 配置中的 unique_kernel_names

非 Inductor 生成的 Triton 内核

CPU 端事件可能不会出现在跟踪中；自动插入探查器事件的机制目前是在 Inductor 层面实现的，因此绕过 Inductor 的 Triton 内核可能不会出现在跟踪中，除非用户手动注释了它们
内核启动应显示为 cuLaunchKernel 而不是 cudaLaunchKernel（cudaLaunchKernel 通常用于 aten 操作）
GPU 端事件应出现，其名称类似于编写的 triton 内核。

Inductor 生成的 CPU 内核

CPU 端事件不会出现在跟踪中；我们尚未为此添加分析功能。
内核启动和GPU 端事件不存在

非 Triton 内核（即 aten 内核或自定义操作）也应预期有时会出现在跟踪中。有时，Inductor 会回退到原始操作的实现，在这种情况下，您将看到对 aten 操作的调用。

启动开销¶

一个常见问题是 GPU 利用率低下。快速识别此问题的方法是，如果 GPU 上的内核之间存在较大的间隔

Visualization in the chrome://trace viewer, showing large gaps between GPU kernels. This indicates that the model is CPU bound, likely due to overhead during kernel launches.

这通常是 CPU 开销的结果，例如，如果在内核启动之间花费在 CPU 上的时间大于 GPU 处理内核所花费的时间。对于较小的批次大小，这个问题更为常见。

使用 inductor 时，启用 CUDA 图表通常可以帮助在启动开销成为问题时提高性能。