來源：DeepHub IMBA本文約6400字，建議閱讀12分鐘在本文將演示 Pytorch 2.0新功能的使用，以及介紹在使用它時可能遇到的一些問題。PyTorch 2.0 發布也有一段時間了，大家是不是已經開始用了呢？PyTorch 2.0 通過引入 torch.compile，可以顯著提高訓練和推理速度。與 eagerly 模式相反，編譯 API 將模型轉換為中間計算圖（FX graph），然後以某種方式將其編譯為低級計算內核，這樣可以提高運行速度。
對於PyTorch 2.0 而言，你看到的可能是：
「只是用 torch.compile 調用包裝它們就可以提高運行速度」
但是其實有許多因素會干擾計算圖編譯和/或達到所需的性能改進。所以需要調整模型和達到最佳性能可能需要重新設計項目或修改一些編碼習慣。
在本文中，我們將演示這個新功能的使用，以及介紹在使用它時可能遇到的一些問題。我們將分享在調整 torch.compile API 時遇到的問題的幾個例子。這些例子並不全面，再實際運用是很可能會遇到此處未提及的問題，並且還要 torch.compile 仍在積極開發中，還有改進的空間。
Torch 編譯背後有許多創新技術，包括 TorchDynamo、FX Graph、TorchInductor、Triton 等。我們不會在這篇文章中深入探討不同的組件，如果你對這些感興趣，可以查看PyTorch 文檔，裡面介紹的非常詳細。
TensorFlow 與 PyTorch 的兩個不重要的對比
1、在過去，PyTorch 和 TensorFlow 之間有著明顯的區別。PyTorch 使用了 eager execution 模式，TensorFlow 使用了 graph 模式，大家都在各自發展。但後來 TensorFlow 2 引入了eager execution作為默認執行模式，TensorFlow 變得有點像 PyTorch。現在 PyTorch 也引入了自己的graph 模式解決方案，變得有點像 TensorFlow。TensorFlow 與 PyTorch 的競爭仍在繼續，但兩者之間的差異正在慢慢消失。
2、人工智慧開發是一個時髦的行業。但是流行的 AI 模型、模型架構、學習算法、訓練框架等隨時間變化發展的。就論文而言，幾年前我們處理的大部分模型都是用 TensorFlow 編寫的。但是人們經常抱怨高級 model.fit API 限制了他們的開發靈活性，並且graph 模式使他們無法調試。然後就有好多人轉向了PyTorch，他們說，「PyTorch可以以任何想要的方式構建模型並輕鬆調試」。但是更靈活的的自定義操作會導致開發的複雜性，PyTorch Lightening等高級的API的出現就是複製了model.fit API的特性，然後同樣的人又說還有人說「我們必須適應 PyTorch Lightening，我們必須用 torch.compile 加速我們的訓練」。既要靈活，又要簡單是不可能同時實現的。
正文開始
下面開始介紹關於如何使用PyTorch 2編譯API的技巧集合，以及一些你可能面臨的潛在問題。使模型適應PyTorch的graph 模式可能需要付出不小的努力。希望這篇文章能幫助你更好地評估這一努力，並決定採取這一步的最佳方式。
安裝PyTorch2
從PyTorch安裝文檔來看，安裝PyTorch 2似乎與安裝任何其他PyTorch版本沒有什麼不同，但是在實踐中，可能會遇到一些問題。首先，PyTorch 2.0(截至本文時)需要Python 3.8或更高版本。然後就是PyTorch 2包含以前版本中不存在的包依賴項(最明顯的是PyTorch-triton，這是什麼我也不知道，哈)，需要注意可能會會引入新的衝突。
所以如果你對Docker熟悉，建議直接使用容器，這樣會簡單很多。
PyTorch2兼容性
PyTorch2的優點之一是它完全向後兼容，所以我們即使不使用torch.compile，仍然可以使用PyTorch 2.0並從其他新功能和增強中受益。最多就是享受不到速度的提升，但是不會有兼容性的問題。但是如果你想進一步提升速度，那麼請往下看。
簡單例子
讓我們從一個簡單的圖像分類模型的例子開始。在下面的代碼塊中，我們使用timm Python包(版本0.6.12)構建一個基本的Vision Transformer (ViT)模型，並在一個假數據集上訓練它500步（不是輪次）。這裡定義了use_compile標誌來控制是否執行模型編譯(torch.compile)，use_amp來控制是使用自動混合精度(AMP)還是全精度(FP)運行。
import time, os
import torch
from torch.utils.data import Dataset
from timm.models.vision_transformer import VisionTransformer
use_amp = True # toggle to enable/disable amp
use_compile = True # toggle to use eager/graph execution mode
# use a fake dataset (random data)
class FakeDataset(Dataset):
def __len__(self):
return 1000000
def __getitem__(self, index):
rand_image = torch.randn([3, 224, 224], dtype=torch.float32)
label = torch.tensor(data=[index % 1000], dtype=torch.int64)
return rand_image, label
def train():
device = torch.cuda.current_device()
dataset = FakeDataset()
batch_size = 64
# define an image classification model with a ViT backbone
model = VisionTransformer()
if use_compile:
model = torch.compile(model)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size, num_workers=4)
loss_function = torch.nn.CrossEntropyLoss()
t0 = time.perf_counter()
summ = 0
count = 0
for idx, (inputs, target) in enumerate(data_loader, start=1):
inputs = inputs.to(device)
targets = torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
with torch.cuda.amp.autocast(
enabled=use_amp,
dtype=torch.bfloat16
):
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss.backward()
optimizer.step()
batch_time = time.perf_counter() - t0
if idx > 10: # skip first few steps
summ += batch_time
count += 1
t0 = time.perf_counter()
if idx > 500:
break
print(f'average step time: {summ/count}')
if __name__ == '__main__':
train()在下表記錄了比較性能結果。這些結果根據環境不同而有很大的變化，所以及供參考
可以看到，使用AMP(28.6%)比使用FP(4.5%)時，模型編譯帶來的性能提升要明顯得多。這是一個眾所周知的差異。如果你還沒有使用AMP進行訓練，那麼其實對於訓練速度的提升是從FP過渡到AMP，所以先推薦你使用AMP。另外就是性能提升伴隨著GPU內存利用率的非常輕微的增加。
當擴展到多個gpu時，由於在編譯圖上實現分布式訓練的方式，比較性能可能會發生變化。具體細節看官方文檔。
https://pytorch.org/get-started/pytorch-2.0/#distributed
高級選項
compile API包含許多用於控制graph創建的選項，能夠針對特定模型對編譯進行微調，並可能進一步提高性能。下面的代碼塊是官方的函數介紹：
def compile(model: Optional[Callable] = None, *,
fullgraph: builtins.bool = False,
dynamic: builtins.bool = False,
backend: Union[str, Callable] = "inductor",
mode: Union[str, None] = None,
options: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] = None,
disable: builtins.bool = False) -> Callable:
"""
Optimizes given model/function using TorchDynamo and specified backend.
Args:
model (Callable): Module/function to optimize
fullgraph (bool): Whether it is ok to break model into several subgraphs
dynamic (bool): Use dynamic shape tracing
backend (str or Callable): backend to be used
mode (str): Can be either "default", "reduce-overhead" or "max-autotune"
options (dict): A dictionary of options to pass to the backend.
disable (bool): Turn torch.compile() into a no-op for testing
"""mode 編譯模式:允許您在最小化編譯所需的開銷(「reduce-overhead」)和最大化潛在的性能提升(「max-autotune」)之間進行選擇。
下表比較了不同編譯模式下編譯上述ViT模型的結果。
可以看到編譯模式的行為與命名的非常相似，「reduce-overhead」以額外的內存利用為代價減少了編譯時間，「max-autotune」以高編譯時間開銷為代價獲得了最佳性能。
backend 編譯器後端:API使用哪個後端將中間表示(IR)計算圖(FX graph)轉換為低級內核操作。這個選項對於調試graph編譯問題和更好地理解torch.compile的內部非常有用。在大多數情況下，默認的Inductor後端似乎能夠提供最佳的訓練性能結果。有很多後端列表，我們可以使用下面命令查看：
from torch import _dynamo
print(_dynamo.list_backends())我們測試使用nvprims-nvfuser後端，可以獲得比eager模式13%的性能提升(與默認後端28.6%的性能提升相比)。具體區別還是要看Pytorch文檔，我們這裡就不細說了，因為文檔都有。
fullgraph 強制單個圖:這個參數是非常有用，可以確保沒有任何不希望的圖截斷。
dynamic 動態形狀:目前 2.0對具有動態形狀的張量的編譯支持在某種程度上是有限的。編譯具有動態形狀的模型的一個常見解決方案是重新編譯，但會大大增加開銷並大大降低訓練速度。如果您的模型確實包含動態形狀，將動態標誌設置為True將帶來更好的性能，特別是減少重新編譯的次數。
都有什麼是動態形狀呢，最簡單的就是時間序列或文本長度不同，如果不進行對齊操作的話序列長度不同就是動態的形狀。
性能分析
PyTorch Profiler是用來分析PyTorch模型性能的關鍵工具之一，可以評估和分析圖編譯優化訓練步驟的方式。在下面的代碼塊中，我們用profiler生成TensorBoard的結果，來查看訓練的性能：
out_path = os.path.join(os.environ.get('SM_MODEL_DIR','/tmp'),'profile')
from torch.profiler import profile, ProfilerActivity
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(
wait=20,
warmup=5,
active=10,
repeat=1),
on_trace_ready=torch.profiler.tensorboard_trace_handler(
dir_name=out_path)
) as p:
for idx, (inputs, target) in enumerate(data_loader, start=1):
inputs = inputs.to(device)
targets = torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
with torch.cuda.amp.autocast(
enabled=use_amp,
dtype=torch.bfloat16
):
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss.backward()
optimizer.step()
p.step()下圖是從PyTorch Profiler生成的TensorBoard 中截取的。它提供了在上面編譯模型試驗的訓練步驟中在GPU上運行的內核的詳細信息。
我們能夠看到torch.compile 增加了GPU張量核心的利用率(從51%到60%)，並且它引入了使用Triton開發的GPU內核。
調試模型編譯問題
torch.compile 目前處於測試階段，如果你遇到問題，並且幸運的話，會得到一個信息錯誤，我們可以直接搜索解決，或者問問chatgpt。但是如果你不那麼幸運，就需要自己尋找問題的根源。
這裡解決編譯問題的主要資源是 TorchDynamo 故障排除文檔，其中包括調試工具列表並提供診斷錯誤的分步指南。但是目前這些工具和技術似乎更多地針對 PyTorch 開發人員而不是 PyTorch 用戶的。它們也許可以幫助解決導致編譯問題的根本問題，但是非常大的可能是它們實際上跟本沒有任何幫助，那怎麼辦呢？
這裡我們演示一個自行解決問題的過程，按照這樣的思路，可以解決一些問題。
下面是一個簡單的分布式模型，其中包括對 torch.distributed.all_reduce 的調用。模型在 eager 模式下按預期運行，但在graph編譯期間失敗並出現「attribute error」（torch.classes.c10d.ProcessGroup does not have a field with name 『shape』）。我們需要將日誌級別提高到 INFO，然後發現發現錯誤在計算的「第 3 步」中，即 TorchInductor。然後通過驗證「eager」和「aot_eager」後端的編譯是否成功， 最後創建一個最小的代碼示例，使用 PyTorch Minifier 重現失敗。
import os, logging
import torch
from torch import _dynamo
# enable debug prints
torch._dynamo.config.log_level = logging.INFO
torch._dynamo.config.verbose=True
# uncomment to run minifier
# torch._dynamo.config.repro_after="aot"
def build_model():
import torch.nn as nn
import torch.nn.functional as F
class DumbNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(1176, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = torch.flatten(x, 1)
x = self.fc1(x)
with torch.no_grad():
sum_vals = torch.sum(x,0)
# this is the problematic line of code
torch.distributed.all_reduce(sum_vals)
# add noise
x = x + 0.1*sum_vals
return x
net = DumbNet()
return net
def train():
os.environ['MASTER_ADDR'] = os.environ.get('MASTER_ADDR',
'localhost')
os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT',
str(2222))
torch.distributed.init_process_group('nccl', rank=0,
world_size=1)
torch.cuda.set_device(0)
device = torch.cuda.current_device()
model = build_model()
model = torch.compile(model)
# replace with this to verfiy that error is not in TorchDynamo
# model = torch.compile(model, 'eager')
# replace with this to verfiy that error is not in AOTAutograd
# model = torch.compile(model, 'aot_eager')
model.to(device)
rand_image = torch.randn([4, 3, 32, 32], dtype=torch.float32).to(device)
model(rand_image)
if __name__ == '__main__':
train()在這個的示例中，運行生成的 minifier_launcher.py 腳本會導致不同的屬性錯誤（比如Repro』 object has no attribute 『_tensor_constant0』），這個對於我們的演示沒有太大幫助，我們暫時忽略他，這也說明了，torch.compile 還不完善，還需要更大的改進空間，或者說如果解決不要問題，那就別用了，至少「慢」要比不能用好，對吧（而且速度提升也有限）
常見的圖截斷問題
Pytorch eager 模式優勢之一是能夠將純 Pythonic 代碼與 PyTorch 操作交織在一起。但是這種自由在使用 torch.compile 時受到很大限制。因為 Pythonic 操作導致 TorchDynamo 將計算圖拆分為多個組件，從而阻礙了性能提升的潛力。而我們代碼優化的目標是儘可能減少此類圖截斷。最簡單的辦法是用 fullgraph 標誌編譯模型。這楊可以提示刪除導致圖截斷的任何代碼，而且還會告訴我們如何最好地適應PyTorch2的開發習慣。但是要運行分布式代碼，則必須將他設為False，因為當前實現 GPU 之間通信的方式需要圖拆分。我們也可以使用 torch._dynamo.explain 序來分析圖截斷。
以下代碼塊演示了一個簡單模型，在其前向傳遞中有四個潛在的圖截斷，但是這種在使用方式在典型的 PyTorch 模型中並不少見。
import torch
from torch import _dynamo
import numpy as np
def build_model():
import torch.nn as nn
import torch.nn.functional as F
class DumbNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(1176, 10)
self.fc2 = nn.Linear(10, 10)
self.fc3 = nn.Linear(10, 10)
self.fc4 = nn.Linear(10, 10)
self.d = {}
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = torch.flatten(x, 1)
assert torch.all(x >= 0) # graph break
x = self.fc1(x)
self.d['fc1-out'] = x.sum().item() # graph break
x = self.fc2(x)
for k in np.arange(1): # graph break
x = self.fc3(x)
print(x) # graph break
x = self.fc4(x)
return x
net = DumbNet()
return net
def train():
model = build_model()
rand_image = torch.randn([4, 3, 32, 32], dtype=torch.float32)
explanation = torch._dynamo.explain(model, rand_image)
print(explanation)
if __name__ == '__main__':
train()圖截斷不會導致編譯失敗(除非設置了fullgraph標誌)。所以很有可能模型正在編譯和運行，但實際上包含多個圖截斷，這會減慢它的速度。
訓練問題故障排除
在目前來說，使用Pytorch2成功編譯的模型就可以認為是一項值得慶祝的成就，但這並不能保證訓練一定會成功。
在 GPU 上運行的低級內核在 eager 模式和graph模式之間會有所不同。某些高級操作可能會表現出不同的行為。你可能會發現在eager 模式下運行的操作在graph 模式下會失敗（例如torch.argmin）。或者會發現計算中的數值差異會影響訓練。
graph模式下的調試比 eager 模式下的調試困難得多。在 eager 模式下，每一行代碼都是獨立執行的，我們可以在代碼中的任意點放置斷點獲得前張量值。而在graph 模式下，代碼定義的模型在處理之前會經歷多次轉換，設置的斷點可能不會被觸發。
所以可以先使用eager 模式，模型跑通以後，再將torch.compile 分別應用於每個部分，或者通過插入列印和/或 Tensor.numpy 調用來生成圖截斷，這樣我們可能會會成功觸發代碼中的斷點。也就是說如果用torch.compile的話對於開發來說，要耗費更長的時間，所以訓練和開發速度的取捨就要看你自己的選擇了。
但是別忘了我們上面說的你的模型在添加了torch.compile後也不一定能正確運行，這又是一個無形的成本。
在圖中包含損失函數
通過使用torch.compile調用包裝PyTorch模型(或函數)來啟用graph模式。但是損失函數不是編譯調用的一部分，也不是生成圖的一部分。所以損失函數是訓練步驟中相對較小的一部分，如果使用eager 模式運行它不會產生太多開銷。但是如果有一個計算量他別大的損失函數，也是可以通過將其包含在編譯的計算圖中來進一步提高性能的。
在下面的代碼中，我們定義了一個損失函數，用於執行從大型ViT模型(具有24個ViT塊)到較小的ViT模型(具有12個ViT塊)的模型蒸餾。
import torch
from timm.models.vision_transformer import VisionTransformer
class ExpensiveLoss(torch.nn.Module):
def __init__(self):
super(ExpensiveLoss, self).__init__()
self.expert_model = VisionTransformer(depth=24)
if torch.cuda.is_available():
self.expert_model.to(torch.cuda.current_device())
self.mse_loss = torch.nn.MSELoss()
def forward(self, input, outputs):
expert_output = self.expert_model(input)
return self.mse_loss(outputs, expert_output)這是一個比CrossEntropyLoss計算量大得多的損失函數，這裡又2種方法讓他執行的更快，
1、loss函數封裝在torch.compile調用中，如下所示:
loss_function = ExpensiveLoss()
compiled_loss = torch.compile(loss_function)這個方法的缺點是損失函數的編譯圖與模型的編譯圖不相交，但是它的優點非常明顯，就是簡單。
2、創建一個包含模型和損失的包裝器模型來將模型和損失一起編譯，並將結果損失作為輸出返回。
import time, os
import torch
from torch.utils.data import Dataset
from torch import nn
from timm.models.vision_transformer import VisionTransformer
# use a fake dataset (random data)
class FakeDataset(Dataset):
def __len__(self):
return 1000000
def __getitem__(self, index):
rand_image = torch.randn([3, 224, 224], dtype=torch.float32)
label = torch.tensor(data=[index % 1000], dtype=torch.int64)
return rand_image, label
# create a wrapper model for the ViT model and loss
class SuperModel(torch.nn.Module):
def __init__(self):
super(SuperModel, self).__init__()
self.model = VisionTransformer()
self.expert_model = VisionTransformer(depth=24 if torch.cuda.is_available() else 2)
self.mse_loss = torch.nn.MSELoss()
def forward(self, inputs):
outputs = self.model(inputs)
with torch.no_grad():
expert_output = self.expert_model(inputs)
return self.mse_loss(outputs, expert_output)
# a loss that simply passes through the model output
class PassthroughLoss(nn.Module):
def __call__(self, model_output):
return model_output
def train():
device = torch.cuda.current_device()
dataset = FakeDataset()
batch_size = 64
# create and compile the model
model = SuperModel()
model = torch.compile(model)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size, num_workers=4)
loss_function = PassthroughLoss()
t0 = time.perf_counter()
summ = 0
count = 0
for idx, (inputs, target) in enumerate(data_loader, start=1):
inputs = inputs.to(device)
targets = torch.squeeze(target.to(device), -1)
optimizer.zero_grad()
with torch.cuda.amp.autocast(
enabled=True,
dtype=torch.bfloat16
):
outputs = model(inputs)
loss = loss_function(outputs)
loss.backward()
optimizer.step()
batch_time = time.perf_counter() - t0
if idx > 10: # skip first few steps
summ += batch_time
count += 1
t0 = time.perf_counter()
if idx > 500:
break
print(f'average step time: {summ/count}')
if __name__ == '__main__':
train()這種方法的缺點是，當在推理模式下運行模型時，需要從包裝器模型中提取內部的實際模型。
這兩種選項的性能提升幅度大致相同都是8%，也就是說，對loss進行編譯也是優化的一個重要部分。
動態形狀
官方也說了torch.compile對動態形狀的模型的編譯支持是有限的。compile API包含dynamic 參數，用於向編譯器發出信號，但是這種方式對於性能提升幫助的程度是值得懷疑的。如果你正在嘗試編譯和優化動態圖並面臨問題，那麼還是不要使用torch.compile，因為太麻煩了。
總結
PyTorch 2.0編譯模式具有顯著提高訓練和推理速度的潛力，可以顯著節省成本，但是模型實現這一潛力所需的工作量可能會有很大差異。許多公共模型只需要修改一行代碼。而其他模型特別是那些包含非標準操作、動態形狀和/或大量交錯Python代碼的模型，可能得不償失甚至無法進行。但是現在開始修改模型是一個很好的選擇，因為目前來看torch.compile對於PyTorch2來說是一個重要且持續的特性。