使用PyTorch 2.0 加速Hugging Face和TIMM库的模型

PyTorch 2.0引入了**torch.compile()**来加速模型，这篇文章我们将介绍如何使用**torch.compile()**来加速Hugging Face和TIMM库的模型。

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 模型，无需更改代码。” 

Sylvain Gugger 是 transformers 和 accelerate 的主要维护者：“只需添加一行代码，PyTorch 2.0 就可以在训练 Transformers 模型时提供 1.5 到 2.x 的加速。这是引入混合精度训练以来最激动人心的事情！”

本教程将向您展示如何使用这些加速，这样您就可以像我们一样对 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 二进制文件中提供了所有必需的依赖项，您可以使用它们下载：

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

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

以 torch.cos() 和 torch.sin() 为特色的示例，它们是逐点操作的示例，因为它们在向量上逐个元素地进行操作。你可能真正想要使用的一个更著名的逐点运算是类似 torch.relu() 的东西。eager模式下的逐点操作不是最优的，因为每个操作都需要从内存中读取一个张量，进行一些更改，然后写回这些更改。

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

回到我们的示例，我们可以将 2 次读取和 2 次写入变成 1 次读取和 1 次写入，这对于较新的 GPU 来说尤其重要，因为瓶颈是内存带宽（您可以多快地向 GPU 发送数据）而不是计算（您的速度有多快） GPU 可以处理浮点运算）。

PyTorch 2.0 为您做的第二个最重要的优化是 CUDA graphs。CUDA graphs有助于消除从 python 程序启动单个内核的开销。

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

TORCHINDUCTOR_TRACE=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”显式传递编译器的名称，但它不是唯一可用的后端，您可以在 torch._dynamo.list_backends() 中运行以查看可用后端的完整列表。为了好玩，您应该尝试 aot_cudagraphs 或 nvfuser。

现在让我们做一些更有趣的事情，我们的社区经常使用来自 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://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)

如果您从模型和 encoded_input 中删除 to(device="cuda:0") ，那么 PyTorch 2.0 将生成 C++ 内核，这些内核将针对在您的 CPU 上运行进行优化。你可以检查 Triton 或 C++ 内核的 BERT，它们显然比我们上面的三角函数示例更复杂，但如果你了解 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 最终成为我们非常有价值的基准测试工具，确保我们所做的任何优化实际上都有助于加速人们想要运行的模型。

本文翻译自https://pytorch.org/blog/Accelerating-Hugging-Face-and-TIMM-models/