开始指南

在阅读本节之前，请先阅读torch.compiler。

让我们从一个简单的torch.compile示例开始，展示如何使用torch.compile进行推理。此示例演示了torch.cos()和torch.sin()功能，这些是逐点操作符的示例，因为它们在向量的每个元素上分别进行操作。虽然这个示例可能不会显示显著的性能提升，但它应该有助于你理解如何在自己的程序中使用torch.compile。

import torch
def fn(x):
   a = torch.cos(x)
   b = torch.sin(a)
   return b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor)

你可能想要使用的一个更著名的逐点操作是像 torch.relu() 这样的函数。在 eager 模式下，每个逐点操作都需要从内存中读取一个张量，进行一些修改，然后再将这些修改写回内存，因此性能较差。inductor 执行的最重要的优化是融合。例如，在上述例子中，我们可以将 2 次读取（x 和 a）和 2 次写入（a 和 b）合并为一次读取（x）和一次写入（b）。这对于新一代 GPU 来说尤为重要，因为这些 GPU 的瓶颈在于内存带宽，即数据传输速度比浮点运算的速度更加关键。

Inductor 提供的另一个重要优化是自动支持 CUDA 图。CUDA 图可以减少从 Python 程序中单独启动内核带来的开销，这对较新的 GPU 来说尤为重要。

TorchDynamo 支持多种不同的后端，而 TorchInductor 通过生成Triton 内核来工作。让我们将上面的例子保存到一个名为 example.py 的文件中。我们可以通过运行 TORCH_COMPILE_DEBUG=1 python example.py 来查看生成的 Triton 内核代码。随着脚本执行，你将在终端上看到 DEBUG 消息。在日志接近结尾时，你会看到一个包含 torchinductor_<your_username> 文件夹路径的信息。在这个文件夹中，你可以找到包含生成内核代码的 output_code.py 文件，类似于以下内容：

@pointwise(size_hints=[16384], filename=__file__, triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
   xnumel = 10000
   xoffset = tl.program_id(0) * XBLOCK
   xindex = xoffset + tl.arange(0, XBLOCK)[:]
   xmask = xindex < xnumel
   x0 = xindex
   tmp0 = tl.load(in_ptr0 + (x0), xmask, other=0.0)
   tmp1 = tl.cos(tmp0)
   tmp2 = tl.sin(tmp1)
   tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

你可以验证 cos 和 sin 确实被融合了，因为这两个操作发生在单个 Triton 内核中，并且临时变量存储在访问速度非常快的寄存器中。

更多关于 Triton 性能的信息可以在 这里 查阅。由于代码是用 Python 编写的，因此即使你不熟悉 CUDA 内核的编写，也能轻松理解。

接下来，让我们尝试使用来自PyTorch_hub的真实模型，例如resnet50。

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(1,3,64,64))

这并非是唯一的可用后端，你可以在REPL中运行torch.compiler.list_backends()来查看所有可用的后端。接下来可以试试cudagraphs作为参考。

import torch
from transformers import BertTokenizer, BertModel
# Copy pasted from here https://huggingface.co/bert-base-uncased
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
model = torch.compile(model, backend="inductor") # This is the only line of code that we changed
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)

import timm
import torch
model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(64,3,7,7))