團隊之所以稱它為 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）恢復性能）。