跳轉到主要內容
部落格

使用 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://huggingface.tw/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。

畢竟，除非您的模型確實執行得更快，否則我們不能聲稱我們建立了一個廣度優先的編譯器。