博客

使用 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 安装了诸如 transformers、accelerate 和 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()，它们是逐点算子（pointwise ops）的示例，即它们对向量中的每个元素进行操作。您可能更想使用的另一个著名的逐点算子是 torch.relu()。

在即时（eager）模式下，逐点算子效率低下，因为每一个算子都需要从内存中读取张量，进行一些更改，然后再写回这些更改。

PyTorch 2.0 为您所做的最重要的优化就是算子融合（fusion）。

回到我们的例子，我们可以将 2 次读取和 2 次写入减少为 1 次读取和 1 次写入。这对于较新的 GPU 尤为重要，因为它们的瓶颈在于内存带宽（即向 GPU 发送数据的速度），而不是计算能力（即 GPU 处理浮点运算的速度）。

PyTorch 2.0 为您提供的第二大优化是 CUDA 图（CUDA graphs）。

CUDA 图有助于消除从 Python 程序启动单个内核（kernel）带来的开销。

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，同样可以快速浏览并理解它们。

如果与 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

毕竟，除非您的模型运行得更快，否则我们不能声称实现了广度优先的覆盖。