跳转到主要内容
博客

使用 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() 这样的。

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

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)

您可以验证两个 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，享受免费的性能提升，如果您没有看到提升，请提出问题，我们将确保您的模型得到支持 https://github.com/pytorch/torchdynamo/issues。

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