作者：Mark Saroufim

torch.compile() 使得通过一行装饰器 torch.compile() 即可轻松尝试不同的编译器后端来加速 PyTorch 代码。它可以直接应用于 nn.Module，作为 torch.jit.script() 的直接替代品，无需修改任何源代码。我们期望这一行代码的修改能够为您已经在运行的绝大多数模型带来 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 就能在训练 Transformers 模型时带来 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 每夜构建（nightly binaries）中提供了所有必需的依赖项，您可以通过以下方式下载

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()，这些都是逐点操作（pointwise ops）的示例，即它们对向量中的每个元素独立进行操作。您实际可能想使用的一个更著名的逐点操作是 torch.relu()。

在 eager 模式下，逐点操作不是最优的，因为每个操作都需要从内存中读取一个张量，进行一些更改，然后将更改写回。

PyTorch 2.0 为您带来的最重要优化是融合（fusion）。

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

PyTorch 2.0 为您带来的第二重要优化是 CUDA graphs。

CUDA graphs 有助于消除从 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)

您可以通过检查来验证这两个 sin 操作确实发生了融合，因为这两个 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 hub 下载一个预训练模型并对其进行优化


import torch
from transformers import BertTokenizer, BertModel
# Copy pasted from here https://hugging-face.cn/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)

如果您从模型和 encoded_input 中移除 to(device="cuda:0")，那么 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

毕竟，除非您的模型真正运行得更快，否则我们不能声称我们创建了一个广度优先的编译器。