PyTorch 2.0 重磅發佈:編譯、編譯、還是編譯!

內容一覽:昨晚召開的 PyTorch Conference 2022 中,官方正式發佈了 PyTorch 2.0。本文將梳理 PyTorch 2.0 與 1.x 相比的最大差異。

關鍵詞:PyTorch 2.0 編譯器 機器學習

在 PyTorch Conference 2022 上,PyTorch 官方正式發佈了 PyTorch 2.0,整場活動含「compiler」率極高,跟先前的 1.x 版本相比,2.0 中有了「顛覆式」的變化。

PyTorch 2.0 中發佈了大量足以改變 PyTorch 使用方式的新功能,它提供了相同的 eager mode 和使用者體驗,同時通過 torch.compile 增加了一個編譯模式,在訓練和推理過程中可以對模型進行加速,從而提供更佳的性能和對 Dynamic Shapes 及 Distributed 的支持。

本文將對 PyTorch 2.0 進行詳細簡介。

太長不看版

* PyTorch 2.0 在保留原有優勢的同時,大舉支持編譯

*torch.compile 為可選功能,只需一行程式碼即可運行編譯

* 4 項重要技術:TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor

* 5 年前就嘗試過編譯,效果一直不理想

* PyTorch 1.x 程式碼無需向 2.0 遷移

* PyTorch 2.0 stable 版本預計明年 3 月發佈

更快、更優、編譯支持

在昨晚的 PyTorch Conference 2022 中,官方正式發佈了 torch.compile,它使得 PyTorch 的性能進一步提升,並開始將 PyTorch 的部分內容從 C++ 中回到 Python。

PyTorch 2.0 中的最新技術包括:

TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor。

1. TorchDynamo

它可以藉助 Python Frame Evaluation Hooks,安全地獲取 PyTorch 程序,這項重大創新是 PyTorch 過去 5 年來在安全圖結構捕獲 (safe graph capture) 方面的研發成果彙總。

2. AOTAutograd

過載 PyTorch autograd engine,作為一個 tracing autodiff,用於生成超前的 backward trace。

3. PrimTorch

將 2000+ PyTorch 運算元歸納為約 250 個 primitive operator 閉集 (closed set),開發者可以針對這些運算元構建一個完整的 PyTorch 後端。PrimTorch 大大簡化了編寫 PyTorch 功能或後端的流程。

4. TorchInductor

一個深度學習編譯器,可以為多個加速器和後端生成 fast code。對於 NVIDIA GPU,它使用 OpenAI Triton 作為關鍵構建模組。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 編寫的,並且支持 dynamic shape(無需重新編譯就能發送不同大小的向量),這使得它們靈活且易學,降低了開發者和供應商的准入門檻。

為了驗證這些技術,PyTorch 官方使用了機器學習領域的 163 個開源模型,包括圖像分類、目標檢測、圖像生成等任務,以及各種 NLP 任務,如語言建模、問答、序列分類、推薦系統和強化學習。這些 Benchmark 分為三類:

* 來自 HuggingFace Transformers 的 46 個模型

* 來自 TIMM 的 61 個模型:由 Ross Wightman 收集的 SoTA PyTorch 圖像模型

* 來自 TorchBench 的 56 個模型:GitHub 上收集的一組流行程式碼庫。

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

接下來 PyTorch 工程師在這些模型中測量速度並驗證精度,由於提速可能取決於資料類型,因此官方在 float32 和自動混合精度 (AMP) 上都測量了提速。由於 AMP 在實踐中更常見,測試比例設定為:0.75 * AMP + 0.25 * float32 的。

在這 163 個開源模型中,torch.compile 可以在 93% 模型上正常運行,運行過後,模型在 NVIDIA A100 GPU 上的運行速度達到了 43% 的提升。在 Float32 精度下,運行速度平均提升 21%;在 AMP 精度下,運行速度平均提升 51%。

注意:在桌面級 GPU(如 NVIDIA 3090)上,測得的速度比在伺服器級 GPU(如 A100)上要低。截至目前,PyTorch 2.0 默認後端 TorchInductor 已經支持 CPU 和 NVIDIA Volta 和 Ampere GPU,暫不支持其他 GPU、xPU 或更老的 NVIDIA GPU。

NVIDIA A100 GPU eager mode

torch.compile 針對不同模型的提速表現

torch.compile 線上試用:開發者可以通過 nightly 二進位制檔案安裝試用,PyTorch 2.0 Stable 版本預計將於 2023 年 3 月初發布。

在 PyTorch 2.x roadmap 中,編譯模式 (compiled mode) 的性能和可擴展性在未來會不斷進行豐富和提升。

PyTorch 2.x roadmap

PyTorch 2.x roadmap

開發背景

PyTorch 的開發理念自始至終都是 flexibility 和 hackability 第一,性能則是第二,致力於:

1. 高性能的 eager execution

2. 不斷 Python 化內部結構

3. Distributed, Autodiff, Data loading, Accelerators 等的良好抽象

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

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

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

PyTorch技術細節

自面世以來,PyTorch 中建立過好幾個編譯器項目,這些編譯器可以分為 3 類:

* 圖結構的獲取 (graph acquisition)

* 圖結構的降低 (graph lowering)

* 圖結構的編譯 (graph compilation)

其中,圖結構的獲取面臨的挑戰最多。

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

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

PyTorch 編譯流程示意圖

PyTorch 編譯流程示意圖

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

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

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

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

越來越多的開發者在編寫高性能自定義核心時,會使用 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)。如果能用上經過驗證的 PyTorch autograd system 就更好了。

AOTAutograd 利用 PyTorch torch_dispatch 擴展機制來追蹤 Autograd engine,使開發者得以「ahead-of-time」捕獲反向傳播 (backwards pas),從而使開發者得以使用 TorchInductor 加速 forwards 和 backwards pass。

PrimTorch:穩定的 Primitive operator

為 PyTorch 寫一個後端並不容易,PyTorch 有 1200+ 運算元,如果考慮到每個運算元的各種過載 (overload),數量高達 2000+。

2000+ PyTorch 運算元的分類概況

因此,編寫後端或交叉功能 (cross-cutting feature) 成為一項耗費精力的工作。PrimTorch 致力於定義更小更穩定的運算元集。PyTorch 程序可以持續降級 (lower) 到這些運算元集。官方的目標是定義兩個運算元集:

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

* ATen ops 包含約 750 個典型運算元 (canonical operator),適合於直接輸出。這些運算元適用於已經在 ATen 級別上集成的後端,或者沒有經過編譯的後端,才能從底層運算元集(如 Prim ops) 恢復性能。

FAQ

1. 如何安裝 PyTorch 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

2. PyTorch 2.0 的程式碼是否向下兼容 1.x?

是的,2.0 不要求修改 PyTorch workflow,只需一行程式碼model = torch.compile(model)即可最佳化模型使用 2.0 stack,並與 PyTorch 其他程式碼順利運行。該選項不強制,開發者仍可使用先前的版本。

3. PyTorch 2.0 是否默認啟用?

不是,必須在 PyTorch 程式碼中明確啟用 2.0,方法是通過一個單一函數調用 (single function call) 來最佳化模型。

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

先前的程式碼不需要任何遷移,如果想使用 2.0 中引入的全新的 compiled mode 功能,可以先用一行程式碼來最佳化模型:model = torch.compile(model)

速度提升主要體現在訓練過程中,如果模型運行速度快於 eager mode,則表示可以用於推理。

import torch

def 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)

5. PyTorch 2.0 中哪些功能不推薦使用 ?

目前 PyTorch 2.0 尚未穩定,且仍處於 nightlies 版本。torch.compile 中的對 Dynamic shapes 的支持仍處於早期,在 2023 年 3 月穩定版 2.0 發佈前不建議使用。

也就是說,即使是 static-shaped workload,仍然在構建 compiled mode,可能會出現一些 bug。對於程式碼中出現崩潰的部分,請禁用 compiled mode 並提交一個 issue。

提交 issue 傳送門:

https://github.com/pytorch/pytorch/issues

以上就是關於PyTorch2.0的詳細介紹,後續我們將完整整理 PyTorch 2.0 Get Started 介紹,歡迎持續關注我們!

您也可以掃描下方二維碼,加入 PyTorch 討論組。

相關文章

深度學習的可解釋性

深度學習的可解釋性

一、深度學習的可解釋性研究概述 隨著深度學習模型在人們日常生活中的許多場景下扮演著越來越重要的角色,模型的「可解釋性」成為了決定使用者是否能...