转摘 PyTorch 2.0 实操：为 HuggingFace 和 TIMM 模型提速！

隋春兰阅读量 14

PyTorch 2.0 通过简单一行 torch.compile() 就可以使模型训练速度提高 30%-200%，本教程将演示如何真实复现这种提速。

torch.compile() 可以轻松地尝试不同的编译器后端， 进而加速 PyTorch 代码的运行。它作为 torch.jit.script() 的直接替代品，可以直接在 nn.Module 上运行，无需修改源代码。

上篇文章中，我们介绍了 torch.compile 支持任意的 PyTorch 代码、control flow、mutation，并一定程度上支持 dynamic shapes。

通过对 163 个开源模型进行测试，我们发现 torch.compile() 可以带来 30%-200% 的加速。

hljs language-ini 复制代码 
opt_module = torch.compile(module)

测试结果详见：

[github.com/pytorch/tor...](https://link.juejin.cn?target=https://github.com/pytorch/torchdynamo/issues/681 "https://github.com/pytorch/torchdynamo/issues/681")

本教程将演示如何利用 torch.compile() 为模型训练提速。

要求及设置

对于 GPU 而言（越新的 GPU 性能提升越突出）：

hljs language-css 复制代码 
pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

对于 CPU 而言：

hljs language-css 复制代码 
pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

可选：验证安装

hljs language-bash 复制代码 
git clone https://github.com/pytorch/pytorchcd tools/dynamopython verify_dynamo.py

可选：Docker 安装

在 PyTorch 的 Nightly Binaries 文件中提供了所有必要的依赖项，可以通过以下方式下载：

hljs language-bash 复制代码 
docker pull ghcr.io/pytorch/pytorch-nightly

对于临时测试 (ad hoc experiment)， 只需确保容器能够访问所有 GPU 即可：

hljs language-bash 复制代码 
docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

开始

简单示例

先来看一个简单示例，注意，GPU 越新速度提升越明显。

hljs language-php 复制代码 
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() 是逐点运算 (pointwise ops) 的例子，他们可以在向量上逐一操作 element，一个更著名的逐点运算是 torch.relu()。

eager mode 下的逐点运算并不是最优解，因为每个算子都需要从内存中读取一个张量、做一些更改，然后再写回这些更改。

PyTorch 2.0 最重要的一项优化是融合 (fusion)。

因此，该例中就可以把 2 次读和 2 次写变成 1 次读和 1 次写，这对较新的 GPU 来说是至关重要的，因为这些 GPU 的瓶颈是内存带宽（能多快地把数据发送到 GPU）而不是计算（GPU 能多快地进行浮点运算）。

PyTorch 2.0 第二个重要优化是 CUDA graphs。

CUDA graphs 有助于消除从 Python 程序中启动单个内核的开销。

torch.compile() 支持许多不同的后端，其中最值得关注的是 Inductor，它可以生成 Triton 内核。

[github.com/openai/trit...](https://link.juejin.cn?target=https://github.com/openai/triton "https://github.com/openai/triton")

这些内核是用 Python 写的，但却优于绝大多数手写的 CUDA 内核。 假设上面的例子叫做 trig.py，实际上可以通过运行来检查生成 triton 内核的代码。

hljs language-python 复制代码 
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 内核中，而且临时变量被保存在 register 中，访问速度非常快。

真实模型示例

以 PyTorch Hub 中的 resnet50 为例：

hljs language-ini 复制代码 
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))

实际运行中会发现，第一次运行速度很慢，这是因为模型正在被编译。随后的运行速度会加快，所以在开始基准测试之前，通常的做法是对模型进行 warm up。

可以看到，这里我们用「inductor」表示编译器名称，但它不是唯一可用的后端，可以在 REPL 中运行 torch._dynamo.list_backends() 来查看可用后端的完整列表。

也可以试试 aot_cudagraphs 或 nvfuser 。

Hugging Face 模型示例

PyTorch 社区经常使用 transformers 或 TIMM 的预训练模型：

[github.com/huggingface...](https://link.juejin.cn?target=https://github.com/huggingface/transformers "https://github.com/huggingface/transformers")

[github.com/rwightman/p...](https://link.juejin.cn?target=https://github.com/rwightman/pytorch-image-models "https://github.com/rwightman/pytorch-image-models")

PyTorch 2.0 的设计目标之一，就是任意编译栈，都需要在实际运行的绝大多数模型中，开箱即用。

这里我们直接从 HuggingFace hub 下载一个预训练的模型，并进行优化：

hljs language-css 复制代码 
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)

如果从模型中删除 to(device="cuda:0") 和 encoded_input ，PyTorch 2.0 将生成为在 CPU 上运行优化的 C++ 内核。

可以检查 BERT 的 Triton 或 C++ 内核，它们显然比上面的三角函数的例子更复杂。但如果你了解 PyTorch 可以略过。

同样的代码与以下一起使用，仍旧可以得到更好的效果：

  1. [github.com/huggingface...](https://link.juejin.cn?target=https://github.com/huggingface/accelerate "https://github.com/huggingface/accelerate")
  2. DDP

同样的，试试 TIMM 的例子：

hljs language-python 复制代码 
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 的目标是建立一个能适配更多模型的编译器，为绝大多数开源模型的运行提速， 现在就访问 HuggingFace Hub，用 PyTorch 2.0 为 TIMM 模型加速吧！

[huggingface.co/timm](https://link.juejin.cn?target=https://huggingface.co/timm "https://huggingface.co/timm")

复制代码 
    ===========================
    【来源： 稀土掘金】
    【作者：  神经星星 】
    【原文链接】 https://juejin.cn/post/7176459007769444410
    声明：转载此文是出于传递更多信息之目的。若有来源标注错误或侵犯了您的合法权益，请作者持权属证明与本网联系，我们将及时更正、删除，谢谢。
标签： PyTorch机器学习人工智能
0/300
全部评论0
0/300