終結1.x時代,PyTorch 2.0來了!100%向後兼容,一行程式碼將訓練提速76%

機器之心編輯部

PyTorch 官方:我們這次的新特性太好用了,所以就直接叫 2.0 了。

前段時間,PyTorch 團隊在官方部落格宣佈 Pytorch 1.13 發佈,包含 BetterTransformer 穩定版等多項更新。在體驗新特性的同時,不少人也在期待下一個版本的推出。

出乎意料的是,這個新版本很快就來了,而且是跨越式的 2.0 版!

新版本的重要進步體現在速度和可用性,而且完全向後兼容。PyTorch 團隊表示,PyTorch 2.0 是他們向 2.x 系列邁出的第一步,其穩定版預計在 2023 年 3 月初發布。

首先,PyTorch 2.0 引入了 torch.compile,這是一種編譯模式,可以在不更改模型程式碼的情況下加速模型。在 163 個涵蓋視覺、NLP 和其他領域的開源模型中,該團隊發現使用 2.0 可以將訓練速度提高 38-76%。

其次,PyTorch 2.0 是 100% 向後兼容的:程式碼庫一樣,API 一樣,寫模型的方式也一樣。團隊之所以稱它為 2.0,是因為它有一些標誌性的新特性,包括:

  • TorchDynamo 可以從位元組碼分析生成 FX 圖;
  • AOTAutograd 可以以 ahead-of-time 的方式生成反向圖;
  • PrimTorch 引入了一個小型運算元集,使後端更容易;
  • TorchInductor:一個由 OpenAI Triton 支持的 DL 編譯器。

PyTorch 2.0 將延續 PyTorch 一貫的優勢,包括 Python 集成、命令式風格、API 簡單等等。此外,PyTorch 2.0 提供了相同的 eager-mode 開發和使用者體驗,同時從根本上改變和增強了 PyTorch 在編譯器級別的運行方式。該版本能夠為「Dynamic Shapes」和分散式運行提供更快的性能和更好的支持。

在官方部落格中,PyTorch團隊還公佈了他們對於整個2.0系列的展望:

以下是詳細內容

以下是詳細內容。

PyTorch 2.X:速度更快、更加地 Python 化、一如既往地 dynamic

PyTorch 2.0 官宣了一個重要特性——torch.compile,這一特性將 PyTorch 的性能推向了新的高度,並將 PyTorch 的部分內容從 C++ 移回 Python。torch.compile 是一個完全附加的(可選的)特性,因此 PyTorch 2.0 是 100% 向後兼容的。

支撐 torch.compile 的技術包括研發團隊新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

  • TorchDynamo 使用 Python Frame Evaluation Hooks 安全地捕獲 PyTorch 程序,這是一項重大創新,是研究團隊對快速可靠地獲取圖進行 5 年研發的結果;
  • AOTAutograd 過載 PyTorch 的 autograd 引擎作為一個跟蹤 autodiff,用於生成 ahead-of-time 向後跟蹤;
  • PrimTorch 將約 2000 多個 PyTorch 運算元規範化為一組約 250 個原始運算元的閉集,開發人員可以將其作為構建完整 PyTorch 後端的目標。這大大降低了編寫 PyTorch 特性或後端的障礙;
  • TorchInductor 是一種深度學習編譯器,可為多個加速器和後端生成快速程式碼。對於 NVIDIA GPU,它使用 OpenAI Triton 作為關鍵構建塊。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 編寫的,並支持 dynamic shapes(即能夠發送不同大小的張量而無需重新編譯),這使得它們具備靈活、易於破解的特性,降低了開發人員和供應商的使用門檻。

為了驗證這些技術,研發團隊在各種機器學習領域測試了 163 個開源模型。實驗精心構建了測試基準,包括各種 CV 任務(圖像分類、目標檢測、圖像生成等)、NLP 任務(語言建模、問答、序列分類、推薦系統等)和強化學習任務,測試模型主要有 3 個來源:

  • 46 個來自 HuggingFace Transformers 的模型;
  • 來自 TIMM 的 61 個模型:一系列 SOTA PyTorch 圖像模型;
  • 來自 TorchBench 的 56 個模型:包含來自 github 的精選流行程式碼庫。

然後研究者測量加速性能並驗證這些模型的準確性。加速可能取決於資料類型,研究團隊選擇測量 float32 和自動混合精度 (AMP) 的加速。

在 163 個開源模型中,torch.compile 在 93% 的情況下都有效,模型在 NVIDIA A100 GPU 上的訓練速度提高了 43%。在 float32 精度下,它的平均運行速度提高了 21%,而在 AMP 精度下,它的運行速度平均提高了 51%。

目前,torch.compile 還處於早期開發階段,預計 2023 年 3 月上旬將發佈第一個穩定的 2.0 版本。

TorchDynamo:快速可靠地獲取圖

TorchDynamo 是一種使用 Frame Evaluation API (PEP-0523 中引入的一種 CPython 特性)的新方法。研發團隊採用資料驅動的方法來驗證其在 Graph Capture 上的有效性,並使用 7000 多個用 PyTorch 編寫的 Github 項目作為驗證集。TorchScript 等方法大約在 50% 的時間裡都難以獲取圖,而且通常開銷很大;而 TorchDynamo 在 99% 的時間裡都能獲取圖,方法正確、安全且開銷可忽略不計(無需對原始程式碼進行任何更改)。這說明 TorchDynamo 突破了多年來模型權衡靈活性和速度的瓶頸。

TorchInductor:使用 define-by-run IR 快速生成程式碼

對於 PyTorch 2.0 的新編譯器後端,研發團隊從使用者編寫高性能自定義核心的方式中汲取靈感:越來越多地使用 Triton 語言。此外,研究者還想要一個編譯器後端——使用與 PyTorch eager 類似的抽象,並且具有足夠的通用性以支持 PyTorch 中廣泛的功能。

TorchInductor 使用 pythonic define-by-run loop level IR 自動將 PyTorch 模型對映到 GPU 上生成的 Triton 程式碼和 CPU 上的 C++/OpenMP。TorchInductor 的 core loop level IR 僅包含約 50 個運算元,並且是用 Python 實現的,易於破解和擴展。

AOTAutograd:將 Autograd 重用於 ahead-of-time 圖

PyTorch 2.0 的主要特性之一是加速訓練,因此 PyTorch 2.0 不僅要捕獲使用者級程式碼,還要捕獲反向傳播。此外,研發團隊還想要複用現有的經過實踐檢驗的 PyTorch autograd 系統。AOTAutograd 利用 PyTorch 的 torch_dispatch 可擴展機制來跟蹤 Autograd 引擎,使其能夠「ahead-of-time」捕獲反向傳遞(backwards pass)。這使 TorchInductor 能夠加速前向和反向傳遞。

PrimTorch:穩定的原始運算元

為 PyTorch 編寫後端具有挑戰性。PyTorch 有 1200 多個運算元,如果考慮每個運算元的各種過載,則有 2000 多個。

在 PrimTorch 項目中,研發團隊致力於定義更小且穩定的運算元集,將 PyTorch 程序縮減到這樣較小的運算元集。目標是定義兩個運算元集:

  • Prim ops:約有 250 個相當低級的運算元。這些運算元適用於編譯器,需要將它們重新融合在一起以獲得良好的性能;
  • ATen ops:約有 750 個規範運算元。這些運算元適用於已經在 ATen 級別集成的後端或沒有編譯功能的後端(無法從較低級別的運算元集(如 Prim ops)恢復性能)。

使用者體驗

PyTorch 2.0 引入了一個簡單的函數 torch.compile,它會返回一個編譯後的模型。

compiled_model *=* torch.compile(model)

compiled_model 保存對模型的引用,並將 forward 函數編譯為一個更最佳化的版本。在編譯模型時,PyTorch 2.0 給了幾項設置來調整它:

def torch.compile(model: Callable,*,mode: Optional[str] = "default",dynamic: bool = False,fullgraph:bool = False,backend: Union[str, Callable] = "inductor",# advanced backend options go here as kwargs**kwargs) -> torch._dynamo.NNOptimizedModule
  • 「mode」指定編譯器在編譯時應該最佳化的內容。
  • default 是一種預設模式,它試圖在不花費太長時間或使用額外記憶體的情況下高效編譯。
  • 其他模式,如 reduce-overhead,可以大大降低框架開銷,但要消耗少量額外記憶體。max-autotune 編譯很長時間,試圖為你提供它所能生成的最快的程式碼。
  • 「dynamic」模式指定是否為 Dynamic Shapes 啟用程式碼路徑。某些編譯器最佳化不能應用於動態形狀的程序。明確你想要一個帶有動態形狀還是靜態形狀的編譯程序,將有助於編譯器提供更好的最佳化程式碼。
  • 「fullgraph」類似於 Numba 的 nopython。它將整個程序編譯成一個圖,或者給出一個錯誤提示,解釋為什麼它不能這樣做。大多數使用者不需要使用這種模式。如果你非常注重性能,那麼你可以嘗試使用它。
  • 「backend 」指定使用哪個編譯器後端。默認情況下使用 TorchInductor,但還有其他一些可用的工具。

編譯體驗想要在默認模式中提供最大的好處和最大的靈活性。

常見問答

1、什麼是 PT 2.0?

2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供相同的 eager 模式開發體驗,同時通過 torch.compile 添加編譯模式。這種編譯模式有可能在訓練和推理期間加速模型。

2、為什麼是 2.0 而不是 1.14?

PyTorch 2.0 就是計劃中的 1.14 。我們發佈了大量新特性,我們相信這些新特性會改變使用者使用 PyTorch 的方式,因此稱其為 PyTorch 2.0。

3、如何安裝 2.0?有什麼額外要求嗎?

安裝最新的 nightlies:

CUDA 11.7

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CUDA 11.6

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116

CPU

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu

4、2.0 程式碼是否向後兼容 1.X?

是的,使用 2.0 不需要修改 PyTorch 工作流程。一行程式碼 model = torch.compile(model) 就可以最佳化模型,以使用 2.0 堆疊,並與其餘 PyTorch 程式碼一起順利運行。這是一個可選擇項,不需要使用新的編譯器。

5、默認啟用 2.0 嗎?

不,必須通過使用單個函數調用最佳化模型從而在 PyTorch 程式碼中顯式啟用 2.0。

6、如何將 PT1.X 程式碼遷移到 PT2.0?

程式碼應該按原樣工作,無需任何遷移。如果想使用 2.0 中引入的新編譯模式特性,那麼可以從最佳化模型開始:

model = torch.compile(model)

雖然加速主要是在訓練期間觀察到的,但如果你的模型運行速度比 eager 模式快,也可以將它用於推理。

import torchdef train(model, dataloader):model = torch.compile(model)for batch in dataloader:run_epoch(model, batch)def infer(model, input):model = torch.compile(model)return model(\*\*input)

7、是否有任何不該使用 PT 2.0 的應用程序?

當前版本的 PT 2.0 仍處於實驗階段,並且處於 nightlies 版本。其中的動態形狀支持還處於早期階段,所以可以等到 2023 年 3 月穩定版發佈後再使用該功能。

8、運行 PyTorch 2.0 時,程式碼有何不同?

開箱即用,PyTorch 2.0 與 PyTorch 1.x 相同,模型以 eager 模式運行,即 Python 的每一行都逐個執行。

在 2.0 中,如果用 model = torch.compile(model) 將模型打包,則模型在執行之前會經過 3 個步驟:

1)圖獲取:首先將模型重寫為子圖塊。可由 TorchDynamo 編譯的子圖被「壓平」,其他子圖(可能包含控制流程式碼或其他不受支持的 Python 結構)將回退到 Eager 模式。

2)Graph lowering:所有 PyTorch 操作都被分解為特定於所選後端的組成核心。

3)圖編譯,核心調用其相應的低級設備專用操作。

9、2.0 目前支持哪些編譯器後端?

默認和最完整的後端是 TorchInductor,但是 TorchDynamo 有一個不斷增長的後端列表,可以通過調用 torchdynamo.list_backends(). 找到

10、2.0 版本的分散式訓練能力如何?

Compiled 模式下的 DDP 和 FSDP ,比 FP32 中的 Eager 模式快 15%、AMP 精度快 80%。PT2.0 做了一些額外的最佳化,以確保 DDP 的通訊 – 計算 overlap 與 Dynamo 的部分圖創建良好協作。想要確保使用 static_graph=False 運行 DDP,更多細節參見:https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860

11、為什麼我的程式碼用 2.0 的 Compiled Model 運行變慢?

性能下降最可能的原因是 graph break 太多。例如,類似模型前向 trigger 中的輸出語句這樣的東西會觸發 graph break。詳見:https://pytorch.org/docs/master/dynamo/faq.html#why-am-i-not-seeing-speedups

12、以前運行的程式碼在 2.0 中崩潰了,該如何調試?

參見:https ://pytorch.org/docs/master/dynamo/faq.html#why-is-my-code-crashing

參考連結:https://pytorch.org/get-started/pytorch-2.0/

相關文章