PyTorch 2.0 釋出，一行程式碼將訓練提速 76%！

在 PyTorch Conference 2022 上，PyTorch 官方正式釋出了 PyTorch 2。0，與先前的 1。x 版本相比，2。0 中有了跨越式的變化。

譯者 | 劉春霖 責編 | 楊紫豔

出品 | CSDN （ID：csdnnews）

在11月初，PyTorch 團隊在官方部落格宣佈 Pytorch 1。13 釋出。據官方介紹，PyTorch 1。13 中包括了 BetterTransformer 的穩定版，且不再支援 CUDA 10。2 及 11。3，並完成了向 CUDA 11。6 及 11。7 的遷移。此外 Beta 版還增加了對 Apple M1 晶片及 functorch 的支援。令人意外的是，12月2日，PyTorch 2。0 正式釋出！與先前的 1。x 版本相比，2。0 中有了跨越式的變化。PyTorch 2。0 中釋出了大量足以改變 PyTorch 使用方式的新功能，它提供了相同的 eager mode 和使用者體驗，同時透過 torch。compile 增加了一個編譯模式，在訓練和推理過程中可以對模型進行加速，從而提供更佳的效能和對 Dynamic Shapes 及分散式執行的支援。PyTorch 團隊表示，PyTorch 2。0 是他們向 2。x 系列邁出的第一步，其穩定版預計在 2023 年 3 月初發布。

PyTorch 2。x：更快、更 Python！

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 程式，這項重大創新是 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 個開源模型，包括影象分類、目標檢測、影象生成等、NLP 任務，如語言建模、問答、序列分類、推薦系統和強化學習任務，測試模型主要有 3 個來源：

46 個來自 HuggingFace Transformers 的模型；

來自 TIMM 的 61 個模型：由 Ross Wightman 收集的SOTA PyTorch 影象模型；

來自 TorchBench 的 56 個模型：包含來自 Github 上收集的一組流行程式碼庫。

對於開源模型，PyTorch 官方沒有進行修改，只是增加了一個 torch。compile 呼叫來進行封裝

接下來 PyTorch 工程師在這些模型中測量速度並驗證精度，由於加速可能取決於資料型別，因此研究團隊選擇測量 Float32 和自動混合精度 （AMP） 的加速。

在 163 個開源模型中，該團隊發現使用 2。0 可以將訓練速度提高 38-76%。torch。compile 在 93% 的情況下都有效，模型在 NVIDIA A100 GPU 上的訓練速度提高了 43%。在 Float32 精度下，它的平均執行速度提高了 21%，而在 AMP 精度下，它的執行速度平均提高了 51%。

目前，torch。compile 還處於早期開發階段，預計 2023 年 3 月上旬將釋出第一個穩定的 2。0 版本。在官方部落格中，PyTorch 團隊還公佈了他們對於整個 2。0 系列的展望：

開發背景

PyTorch 的開發理念自始至終都是靈活性和 hackability 第一，效能則是第二，致力於：

1。 高效能的 eager execution

2。 不斷 Python 化內部結構

3。 分散式、自動比較、資料載入、加速器等的良好抽象

PyTorch 自 2017 年面世以來，硬體加速器（如 GPU）的計算速度提高了約 15倍，記憶體訪問速度提高了約 2 倍。

為了保持高效能的 eager execution，PyTorch 內部的大部分內容不得不轉移到 C++ 中，這使得 PyTorch hackability 下降，也增加了開發者參與程式碼貢獻的門檻。

從第一天起，PyTorch 官方就意識到了 eager execution 的效能侷限。2017 年 7 月，官方開始致力於為 PyTorch 開發一個編譯器。該編譯器需要在不犧牲 PyTorch 體驗的前提下，加速 PyTorch 程式的執行，其關鍵標準是保持某種程度上的靈活性 （flexibility）：支援開發者廣泛使用的 dynamic shapes 以及 dynamic programs。

開發者 Sylvain Gugger 表示：“只需新增一行程式碼，PyTorch 2。0 就能在訓練 Transformers 模型時實現 1。5 倍到 2。0 倍的速度提升。這是自混合精度訓練問世以來最令人興奮的事情！”

技術概述

多年來，研究者們在 PyTorch 中建立過好幾個編譯器專案，這些編譯器可以分為 3 類：

圖結構的獲取圖結構的降低圖結構的編譯

其中，在構建 PyTorch 編譯器時，圖結構的獲取是更難的挑戰。

過去 5 年中，官方嘗試了 torch。jit。trace、TorchScript、FX tracing 以及 Lazy Tensors，但它們有些夠靈活但不夠快，有些夠快但不靈活，有些既不快也不靈活，有些使用者體驗不好。

雖然 TorchScript 很有前途，但它需要大量修改程式碼和依賴，可行性並不高。

PyTorch 編譯流程示意圖

TorchDynamo：可靠快速地獲取圖結構

TorchDynamo 使用了 PEP-0523 中引入的CPython 功能，稱為框架評估 API （Frame Evaluation API）。為此，官方採取了一種資料驅動的方法來驗證其在 Graph Capture 上的有效性，使用 7000 多個用 PyTorch 編寫的 Github 專案作為驗證集。

結果顯示，TorchDynamo 在 99% 的時間裡都能正確、安全地獲取圖結構，而且開銷可以忽略不計，因為它無需對原始程式碼做任何修改。

TorchInductor：用 define-by-run IR 進行更迅速的 codegen

對於 PyTorch 2。0 的新編譯器後端，團隊從使用者如何編寫高效能的自定義核心中得到了靈感：越來越多地使用 Triton 語言。

此外，對於 PyTorch 2。0 全新的編譯器後端，官方還希望能夠使用與 PyTorch eager 類似的抽象，並且具有足夠的通用效能支援 PyTorch 中廣泛的功能。

TorchInductor 使用 Pythonic define-by-run loop level IR，自動將 PyTorch 模型對映到 GPU 上生成的 Triton 程式碼以及 CPU 上的 C++/OpenMP。

TorchInductor 的 core loop level IR 只包含大約 50 個運算元，而且是用 Python 實現的，這使得它具有很強的 hackability 和擴充套件性。

AOTAutograd：對於 ahead-of-time graph，重用 Autograd

PyTorch 2。0 要想加速訓練，不僅要捕獲使用者級程式碼，而且要捕獲反向傳播演算法 （backpropagation）。

AOTAutograd 利用 PyTorch torch_dispatch 擴充套件機制來追蹤 Autograd 引擎，使開發者得以提前捕獲反向傳播 （backwards pas），從而使開發者得以使用 TorchInductor 加速前向和後向通道。

PrimTorch：穩定的 Primitive operator

PyTorch 有 1200多個運算子，如果考慮到每個運算子的各種過載，數量高達 2000+。

2000+ PyTorch 運算元的分類概況

因此，編寫後端或交叉功能 （cross-cutting feature） 成為一項耗費精力的工作。PrimTorch 致力於定義更小更穩定的運算子集。PyTorch 程式可以持續降級到這些運算子集。

官方的目標是定義兩個運算子集：

Prim ops 包含約 250 個相對底層的運算子，因為足夠底層，所以這些運算子更適用於編譯器，開發者需要將這些運算元進行融合，才能獲得良好的效能。

ATen ops 包含約 750 個典型運算子 （canonical operator），適合於直接輸出。這些運算子適用於已經在 ATen 級別上整合的後端，或者沒有經過編譯的後端，才能從底層運算子集（如 Prim ops） 恢復效能。

動態形狀

在研究支援 PyTorch 程式碼通用性的必要條件時，一個關鍵要求是支援動態形狀，並允許模型接受不同大小的張量，而不會在每次形狀變化時引起重新編譯。目前為止，對動態形狀的支援有限，並且正在進行中。它將在穩定版本中具有完整的功能。在不支援動態形狀的情況下，常見的解決方法是將其填充到最接近的 2 次方。然而，正如我們從下面的圖表中所看到的，它產生了大量的效能開銷，同時也帶來了明顯更長的編譯時間。現在，有了對動態形狀的支援，PyTorch 2。0 也就獲得了比 Eager 高出了最多 40% 的效能。

簡而言之，就是 PyTorch 2。0 stable 版本預計明年 3 月釋出，PyTorch 2。0 在保留原有優勢的同時，大力支援編譯，torch。compile 為可選功能，只需一行程式碼即可執行編譯，還有 4 項重要技術：TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor，PyTorch 1。x 程式碼無需向 2。0 遷移，更多使用者體驗以及Q&A，大家可以檢視下方連結中原文呦~

參考連結：

https：//pytorch。org/get-started/pytorch-2。0/

☞☞