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

來源：NG影片 釋出時間：2022-12-07 10:02

在 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+。