跳转到主要内容
博客

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

作者： 2022 年 12 月 2 日2024 年 11 月 14 日暂无评论

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 每夜版二进制文件中提供了所有必需的依赖项，您可以通过以下方式下载

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 为您做的最重要的优化是融合。

所以回到我们的例子，我们可以将两次读取和两次写入转换为一次读取和一次写入，这至关重要，特别是对于较新的 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 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 将生成 C++ 内核，这些内核将针对在您的 CPU 上运行进行优化。您可以检查 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，享受免费的性能提升，如果您没有看到，请提交一个问题，我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues。

毕竟，除非您的模型实际运行速度更快，否则我们无法声称我们创建了广度优先的解决方案。