分析以理解 torch.compile 性能

使用 torch.profiler 的目的：

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

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

另请参阅一般pytorch分析器指南。

使用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()

# 预热
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”：

或者，在顶部使用“Flow events”下拉菜单开启所有流程。

解决CUDA图分析问题

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

import torch

torch.profiler._utils._init_for_cuda_graphs()

# ... 程序的其余部分

理解编译时间

要了解为什么编译需要很长时间，您可以分析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()

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

prof.export_chrome_trace("trace_compile.json")

注意几点：

• 第一次调用应在分析期间进行，以便捕获编译

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

查找图表断点

尽管有用于识别图形中断的日志工具，但探查器提供了一种快速的可视化方法来识别图形中断。

当需要任何输入的梯度时，图断裂很容易识别：每个图断裂都会中断一个CompiledFunction块，将其分成两部分。

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

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()

# 预热
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")

启动开销

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

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

在使用电感器时，启用CUDA图可以经常帮助在启动开销是一个问题时提高性能。