由马克·萨拉菲姆所著

torch.compile() 使您能够轻松尝试不同的编译器后端，通过单行装饰器 torch.compile() . 来加速 PyTorch 代码。它可以直接作为 nn.Module 的替代品使用，而无需对源代码进行任何修改。我们预计这一行代码的更改将为您在大多数已运行的模型上提供 30%-2 倍的训练时间加速。


opt_module = torch.compile(module)

torch.compile 支持任意 PyTorch 代码、控制流、突变，并附带对动态形状的实验性支持。我们对这一发展感到非常兴奋，因此将其称为 PyTorch 2.0。

对于我们来说，这个公告的不同之处在于，我们已经对一些最受欢迎的开源 PyTorch 模型进行了基准测试，并获得了从 30% 到 2 倍的显著加速。https://github.com/pytorch/torchdynamo/issues/681。

这里没有技巧，我们已通过 pip 安装了流行的库，如 https://github.com/huggingface/transformers、https://github.com/huggingface/accelerate 和 https://github.com/rwightman/pytorch-image-models，然后对它们运行了 torch.compile()，就这样了。

同时获得性能和便利性是罕见的，这就是为什么核心团队对 PyTorch 2.0 感到如此兴奋。Hugging Face 团队也感到兴奋，正如他们所说：

TIMM 的主要维护者 Ross Wightman：“PT 2.0 与大多数 timm 模型在推理和训练工作负载上无缝工作，无需代码更改。”

transformers 和 accelerate 的主要维护者 Sylvain Gugger：“只需添加一行代码，PyTorch 2.0 就能在训练 Transformer 模型时将速度提升 1.5 倍到 2 倍。这是自混合精度训练引入以来最令人兴奋的事情！”

本教程将向您展示如何精确复制这些加速，让您像我们一样对 PyTorch 2.0 感到兴奋。

需求和设置

对于 GPU（新一代 GPU 将看到显著更好的性能）

pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

对于 CPU

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

可选：验证安装

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

可选：Docker 安装

我们还在 PyTorch 夜间构建中提供了所有必需的依赖项，您可以通过以下方式下载：

docker pull ghcr.io/pytorch/pytorch-nightly

对于临时实验，请确保您的容器可以访问所有 GPU

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

开始使用

一个玩具示例

让我们从简单的示例开始，逐步使事情变得复杂。请注意，您使用的 GPU 越新，您可能会看到更显著的加速。

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

这个示例实际上不会运行得更快，但它具有教育意义。

这是一种包含 torch.cos() 和 torch.sin() 的示例，它们是逐点操作，即它们逐元素对向量进行操作。你可能真正想使用的更著名的逐点操作可能是类似 torch.relu() 的东西。

在动态计算模式下，逐点操作效率不高，因为每个操作都需要从内存中读取一个张量，进行一些更改，然后再将更改写回。

PyTorch 2.0 为你做的最重要的优化是融合。

因此，回到我们的示例，我们可以将 2 次读取和 2 次写入转换为 1 次读取和 1 次写入，这对于内存带宽（将数据发送到 GPU 的速度）是至关重要的，尤其是在内存带宽成为瓶颈的新一代 GPU 上，而不是计算（GPU 执行浮点运算的速度）。

PyTorch 2.0 为你做的第二个最重要的优化是 CUDA 图

CUDA 图有助于消除从 Python 程序中启动单个内核的开销。

torch.compile()支持许多不同的后端，但我们特别兴奋的是 Inductor，它生成 Triton 内核（https://github.com/openai/triton），这些内核是用 Python 编写的，但性能优于绝大多数手写的 CUDA 内核。假设我们的示例文件名为 trig.py，我们实际上可以通过运行来检查生成的 Triton 内核代码。

TORCH_COMPILE_DEBUG=1 python trig.py


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

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

真实模型

作为下一步，让我们尝试一个真实的模型，比如来自 PyTorch hub 的 resnet50。

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

如果您实际运行，可能会惊讶地发现第一次运行很慢，这是因为模型正在编译。随后的运行将会更快，所以在开始基准测试之前预热模型是一种常见的做法。

您可能已经注意到我们在这里明确传递了编译器的名称“inductor”，但这并不是唯一可用的后端，您可以在 REPL torch._dynamo.list_backends() 中运行以查看所有可用的后端。为了好玩，您应该尝试 aot_cudagraphs 或 nvfuser 。

Hugging Face 模型

现在我们来做点更有趣的事情，我们的社区经常使用来自 transformers https://github.com/huggingface/transformers 或 TIMM https://github.com/rwightman/pytorch-image-models 的预训练模型，并且 PyTorch 2.0 的一个设计目标是任何新的编译器堆栈都需要与人们实际运行的绝大多数模型无缝兼容。

因此，我们将直接从 Hugging Face 网络下载一个预训练模型并进行优化


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

如果你从模型中移除 to(device="cuda:0") 和 encoded_input ，PyTorch 2.0 将生成针对 CPU 运行的 C++ 内核。你可以检查 BERT 的 Triton 或 C++ 内核，它们显然比我们上面提到的三角函数示例复杂，但如果你理解 PyTorch，你可以同样快速浏览并理解。

同样的代码与 https://github.com/huggingface/accelerate 和 DDP 一起使用也运行良好

同样，让我们尝试一个 TIMM 示例

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

我们使用 PyTorch 的目标是构建一个广度优先的编译器，以加快人们在开源中运行的绝大多数实际模型的速度。Hugging Face Hub 最终成为我们极具价值的基准测试工具，确保我们进行的任何优化实际上有助于加速人们想要运行的模型。

因此，请尝试 PyTorch 2.0，享受免费的性能提升，如果您没有看到效果，请提交一个 issue，我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues

毕竟，我们不能声称我们创建了一个广度优先搜索，除非你的模型实际上运行得更快。