團(tuán)隊(duì)之所以稱它為 2.0，是因?yàn)樗幸恍?biāo)志性的新特性，包括：

 

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

 

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

 

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

 

　　TorchInductor：一個(gè)由 OpenAI Triton 支持的 DL 編譯器。

 

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

 

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

 

　

 

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

 

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

 

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

 

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

 

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

 

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

 

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

 

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

 

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

 

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

 

　　46 個(gè)來自 HuggingFace Transformers 的模型；

 

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

 

　　來自 TorchBench 的 56 個(gè)模型：包含來自 github 的精選流行代碼庫(kù)。

 

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

 

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

 

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

 

　　TorchDynamo：快速可靠地獲取圖

 

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

 

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

 

　　對(duì)于 PyTorch 2.0 的新編譯器后端，研發(fā)團(tuán)隊(duì)從用戶編寫高性能自定義內(nèi)核的方式中汲取靈感：越來越多地使用 Triton 語(yǔ)言。此外，研究者還想要一個(gè)編譯器后端——使用與 PyTorch eager 類似的抽象，并且具有足夠的通用性以支持 PyTorch 中廣泛的功能。

 

　　TorchInductor 使用 pythonic define-by-run loop level IR 自動(dòng)將 PyTorch 模型映射到 GPU 上生成的 Triton 代碼和 CPU 上的 C++/OpenMP。TorchInductor 的 core loop level IR 僅包含約 50 個(gè)算子，并且是用 Python 實(shí)現(xiàn)的，易于破解和擴(kuò)展。

 

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

 

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

 

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

 

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

 

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

 

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

 

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