分析以了解 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()
# 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")
注意以下几点
第一次调用应该发生在分析期间，以便捕获编译
添加一个预热编译，以初始化任何需要延迟初始化的系统。
查找图形断点¶
虽然有用于识别图形断点的日志记录工具，但分析器提供了一种快速识别图形断点的可视化方法。
当任何输入都需要梯度时，图形断点很容易识别：每个图形断点都会中断一个 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()
# 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")
启动开销¶
一个常见问题是 GPU 利用率低下。快速识别此问题的方法是，如果 GPU 上的内核之间存在较大的间隙。
这通常是 CPU 开销的结果，例如，如果内核启动之间在 CPU 上花费的时间大于 GPU 处理内核所花费的时间。对于较小的批次大小，这个问题更为常见。
在使用 Inductor 时，启用 CUDA 图表通常可以帮助提高性能，尤其是在启动开销成为问题时。
