torch.compiler 基础

`torch.compiler` 基础#

先来看一个简单的 torch.compile() 进行推理示例。这个例子展示了 torch.cos() 和 torch.sin() 函数，它们是点对点算子的例子，因为它们在向量上逐元素运算。这个例子可能不会显示显著的性能提升，但应该能帮助你形成自己程序中如何使用 torch.compile 的直观理解。

import torch
torch._dynamo.reset()
def fn(x):
   a = torch.cos(x)
   b = torch.sin(a)
   return b

new_fn = torch.compile(fn, backend="cudagraphs")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor)

在 eager 模式下，逐点运算并不理想，因为每个运算都需要从内存中读取一个张量，进行一些更改，然后将这些更改写回。inductor 执行的最主要优化是融合。在上述示例中，我们可以将 2 个读取 (x, a) 和 2 个写入 (a, b) 转换为 1 个读取 (x) 和 1 个写入(b)，这对于新的 GPU 来说至关重要，因为瓶颈在于内存带宽（将数据发送到 GPU 的速度）而不是计算（GPU 可以快速处理浮点运算的速度）。

inductor 还提供了对 CUDA graph 的自动支持。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__, 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)
   tmp1 = tl.cos(tmp0)
   tmp2 = tl.sin(tmp1)
   tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

你可以验证融合 cos 和 sin 的算子确实发生了，因为 cos 和 sin 运算发生在单个 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")
out = opt_model(torch.randn(1,3,64,64))

查看支持的后端：

torch.compiler.list_backends()

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

PyTorch 用户经常利用来自 transformers 或 TIMM 的预训练模型，其中设计目标是让 TorchDynamo 和 TorchInductor 能够与人们想要编写的任何模型一起开箱即用。

让我们直接从 HuggingFace hub 下载预训练模型并进行优化：

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='cudagraphs') # 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)

如果你从模型和 encoded_input 中删除 to(device="cuda:0")，那么 Triton 将生成针对你的 CPU 优化的 C++ 内核。你可以检查 BERT 的 Triton 或 C++ 内核。它们比我们尝试过的三角函数示例更复杂。

同样，让我们尝试 TIMM 示例：

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

torch.compiler 基础

torch.compiler 基础#

`torch.compiler` 基础#