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% 的加速。
opt_module = torch.compile(module)
測試結果詳見：
https://github.com/pytorch/to...
本教程將演示如何利用 torch.compile() 為模型訓練提速。
要求及設定
對於 GPU 而言（越新的 GPU 效能提升越突出）：
pip3 install numpy --pre torch[dynamo] --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 的 Nightly Binaries 檔案中提供了所有必要的依賴項，可以透過以下方式下載：
docker pull ghcr.io/pytorch/pytorch-nightly對於臨時測試 (ad hoc experiment)，只需確保容器能夠訪問所有 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()這個例子實際上不會提升速度，但是可以拋磚引玉。
該示例中，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 核心。
https://github.com/openai/triton
這些核心是用 Python 寫的，但卻優於絕大多數手寫的 CUDA 核心。假設上面的例子叫做 trig.py，實際上可以透過執行來檢查生成 triton 核心的程式碼。
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 為例：
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 的預訓練模型：
https://github.com/huggingfac...
https://github.com/rwightman/...
PyTorch 2.0 的設計目標之一，就是任意編譯棧，都需要在實際執行的絕大多數模型中，開箱即用。
這裡我們直接從 HuggingFace hub 下載一個預訓練的模型，並進行最佳化：
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 可以略過。
同樣的程式碼與以下一起使用，仍舊可以得到更好的效果：
同樣的，試試 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 的目標是建立一個能適配更多模型的編譯器，為絕大多數開源模型的執行提速，現在就訪問 HuggingFace Hub，用 PyTorch 2.0 為 TIMM 模型加速吧！