跳轉到主要內容
部落格

使用 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 將生成針對在 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。

畢竟，除非您的模型實際執行得更快，否則我們不能聲稱我們建立了一個廣度優先的解決方案。