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

torch.profiler 的用途：

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

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

另请参阅通用 pytorch profiler 指南。

使用 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”键左右滚动。 “？”将显示包含快捷键列表的“帮助”屏幕。

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

CPU 和 GPU 事件之间的流

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

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

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

解决 CUDA 图性能分析问题

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

import torch

torch.profiler._utils._init_for_cuda_graphs()

# ... rest of program

理解编译时间

为了理解为什么编译需要很长时间，您可以对 torch.compile-ed 程序的首次调用进行性能分析。 请记住，编译的性能分析跟踪可能会比典型的性能分析更失真，因为编译工作负载可能与典型的 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")

注意以下几点

• 首次调用应在性能分析期间发生，以便捕获编译

• 添加预热编译，以便初始化任何需要延迟初始化的系统。

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

虽然有一些日志记录工具可以识别图中断，但 profiler 提供了一种快速可视化的方法来识别图中断。 有两个 profiler 事件需要注意：Torch-Compiled Region 和 CompiledFunction。

Torch-Compiled Region - 在 PyTorch 2.2 中引入 - 是一个 profiler 事件，它涵盖了整个编译区域。 图中断几乎总是看起来相同：嵌套的“Torch-Compiled Region”事件。

如果您对两个独立的函数应用 torch.compile()，您通常应该看到两个相邻的（即非堆叠/嵌套）Torch-Compiled 区域。 同时，如果您遇到图中断（或 disable()’ed/skipped 区域），则预期会出现嵌套的“Torch-Compiled Region”事件。

CompiledFunction - 在 PyTorch 2.0 中引入 - 是当任何输入需要梯度时出现的 profiler 事件。 每个图中断都会中断 CompiledFunction 块，将其拆分为两个。 CompiledFunction 事件仅在涉及 Autograd 时出现，即图的某些输入张量具有 requires_grad=True。

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

如果您的用例包含不需要梯度且不包含“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")

运算符内核

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

1. CPU 端事件

2. 内核启动（如果处理 GPU 内核）

3. GPU 端事件

Inductor 生成的 Triton 内核： 1. CPU 端事件应显示为前缀为“triton_”的事件。 这些事件目前包含最少的信息 - 内核名称和启动，但信息少于典型的 aten 内核启动（其中包含输入形状、类型等）。 2. 内核启动应显示为 cuLaunchKernel 而不是 cudaLaunchKernel（cudaLaunchKernel 是 aten ops 的典型情况） 3. GPU 端事件应显示，名称的描述性程度取决于 unique_kernel_names 的 inductor 配置

非 Inductor 生成的 Triton 内核

1. CPU 端事件可能不会出现在跟踪中； 用于自动插入 profiler 事件的机制目前在 Inductor 级别实现，因此绕过 Inductor 的 Triton 内核可能不会出现在跟踪中，除非用户手动注释它们

2. 内核启动应显示为 cuLaunchKernel 而不是 cudaLaunchKernel（cudaLaunchKernel 是 aten ops 的典型情况）

3. GPU 端事件应显示，名称与编写的 triton 内核类似。

Inductor 生成的 CPU 内核

1. CPU 端事件不会出现在跟踪中； 我们尚未为此添加性能分析。

2. 内核启动和 GPU 端事件不存在

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

启动开销

一个常见的问题是 GPU 利用率不佳。 快速识别这一点的方法是 GPU 上内核之间存在很大的间隔

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

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