團(tuán)隊之所以稱它為 2.0，是因為它有一些標(biāo)志性的新特性，包括：

 

　　TorchDynamo 可以從字節(jié)碼分析生成 FX 圖；

 

　　AOTAutograd 可以以 ahead-of-time 的方式生成反向圖；

 

　　PrimTorch 引入了一個小型算子集，使后端更容易；

 

　　TorchInductor：一個由 OpenAI Triton 支持的 DL 編譯器。

 

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

 

　　在官方博客中，PyTorch團(tuán)隊還公布了他們對于整個2.0系列的展望：

 

　

 

　　以下是詳細(xì)內(nèi)容。

 

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

 

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

 

　　支撐 torch.compile 的技術(shù)包括研發(fā)團(tuán)隊新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

 

　　TorchDynamo 使用 Python Frame Evaluation Hooks 安全地捕獲 PyTorch 程序，這是一項重大創(chuàng)新，是研究團(tuán)隊對快速可靠地獲取圖進(jìn)行 5 年研發(fā)的結(jié)果；

 

　　AOTAutograd 重載 PyTorch 的 autograd 引擎作為一個跟蹤 autodiff，用于生成 ahead-of-time 向后跟蹤；

 

　　PrimTorch 將約 2000 多個 PyTorch 算子規(guī)范化為一組約 250 個原始算子的閉集，開發(fā)人員可以將其作為構(gòu)建完整 PyTorch 后端的目標(biāo)。這大大降低了編寫 PyTorch 特性或后端的障礙；

 

　　TorchInductor 是一種深度學(xué)習(xí)編譯器，可為多個加速器和后端生成快速代碼。對于 NVIDIA GPU，它使用 OpenAI Triton 作為關(guān)鍵構(gòu)建塊。

 

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

 

　　為了驗證這些技術(shù)，研發(fā)團(tuán)隊在各種機器學(xué)習(xí)領(lǐng)域測試了 163 個開源模型。實驗精心構(gòu)建了測試基準(zhǔn)，包括各種 CV 任務(wù)（圖像分類、目標(biāo)檢測、圖像生成等）、NLP 任務(wù)（語言建模、問答、序列分類、推薦系統(tǒng)等）和強化學(xué)習(xí)任務(wù)，測試模型主要有 3 個來源：

 

　　46 個來自 HuggingFace Transformers 的模型；

 

　　來自 TIMM 的 61 個模型：一系列 SOTA PyTorch 圖像模型；

 

　　來自 TorchBench 的 56 個模型：包含來自 github 的精選流行代碼庫。

 

　　然后研究者測量加速性能并驗證這些模型的準(zhǔn)確性。加速可能取決于數(shù)據(jù)類型，研究團(tuán)隊選擇測量 float32 和自動混合精度 (AMP) 的加速。

 

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

 

　　目前，torch.compile 還處于早期開發(fā)階段，預(yù)計 2023 年 3 月上旬將發(fā)布第一個穩(wěn)定的 2.0 版本。

 

　　TorchDynamo：快速可靠地獲取圖

 

　　TorchDynamo 是一種使用 Frame Evaluation API （PEP-0523 中引入的一種 CPython 特性）的新方法。研發(fā)團(tuán)隊采用數(shù)據(jù)驅(qū)動的方法來驗證其在 Graph Capture 上的有效性，并使用 7000 多個用 PyTorch 編寫的 Github 項目作為驗證集。TorchScript 等方法大約在 50% 的時間里都難以獲取圖，而且通常開銷很大；而 TorchDynamo 在 99% 的時間里都能獲取圖，方法正確、安全且開銷可忽略不計（無需對原始代碼進(jìn)行任何更改）。這說明 TorchDynamo 突破了多年來模型權(quán)衡靈活性和速度的瓶頸。

 

　　TorchInductor：使用 define-by-run IR 快速生成代碼

 

　　對于 PyTorch 2.0 的新編譯器后端，研發(fā)團(tuán)隊從用戶編寫高性能自定義內(nèi)核的方式中汲取靈感：越來越多地使用 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 實現(xiàn)的，易于破解和擴展。

 

　　AOTAutograd：將 Autograd 重用于 ahead-of-time 圖

 

　　PyTorch 2.0 的主要特性之一是加速訓(xùn)練，因此 PyTorch 2.0 不僅要捕獲用戶級代碼，還要捕獲反向傳播。此外，研發(fā)團(tuán)隊還想要復(fù)用現(xiàn)有的經(jīng)過實踐檢驗的 PyTorch autograd 系統(tǒng)。AOTAutograd 利用 PyTorch 的 torch_dispatch 可擴展機制來跟蹤 Autograd 引擎，使其能夠「ahead-of-time」捕獲反向傳遞（backwards pass）。這使 TorchInductor 能夠加速前向和反向傳遞。

 

　　PrimTorch：穩(wěn)定的原始算子

 

　　為 PyTorch 編寫后端具有挑戰(zhàn)性。PyTorch 有 1200 多個算子，如果考慮每個算子的各種重載，則有 2000 多個。

 

　　在 PrimTorch 項目中，研發(fā)團(tuán)隊致力于定義更小且穩(wěn)定的算子集，將 PyTorch 程序縮減到這樣較小的算子集。目標(biāo)是定義兩個算子集：

 

　　Prim ops：約有 250 個相當(dāng)?shù)图壍乃阕?。這些算子適用于編譯器，需要將它們重新融合在一起以獲得良好的性能；

 

　　ATen ops：約有 750 個規(guī)范算子。這些算子適用于已經(jīng)在 ATen 級別集成的后端或沒有編譯功能的后端（無法從較低級別的算子集（如 Prim ops）恢復(fù)性能）。