12月2日，PyTorch 2.0正式發(fā)布！

這次的更新不僅將PyTorch的性能推到了新的高度，同時(shí)也加入了對(duì)動(dòng)態(tài)形狀和分布式的支持。

此外，2.0系列還會(huì)將PyTorch的部分代碼從C++移回Python。

(資料圖)

目前，PyTorch2.0還處在測(cè)試階段，預(yù)計(jì)第一個(gè)穩(wěn)定版本會(huì)在2023年3月初面世。

PyTorch 2.x：更快、更Python！

在過(guò)去的幾年里，PyTorch從1.0到最近的1.13進(jìn)行了創(chuàng)新和迭代，并轉(zhuǎn)移到新成立的PyTorch基金會(huì)，成為L(zhǎng)inux基金會(huì)的一部分。

當(dāng)前版本的PyTorch所面臨的挑戰(zhàn)是，eager-mode難以跟上不斷增長(zhǎng)的GPU帶寬和更瘋狂的模型架構(gòu)。

而PyTorch 2.0的誕生，將從根本上改變和提升了PyTorch在編譯器級(jí)別下的運(yùn)行方式。

眾所周知，PyTorch中的（Py）來(lái)自于數(shù)據(jù)科學(xué)中廣泛使用的開源Python編程語(yǔ)言。

然而，PyTorch的代碼卻并沒(méi)有完全采用Python，而是把一部分交給了C++。

不過(guò)，在今后的2.x系列中，PyTorch項(xiàng)目團(tuán)隊(duì)計(jì)劃將與torch.nn有關(guān)的代碼移回到Python中。

除此之外，由于PyTorch 2.0是一個(gè)完全附加的（和可選的）功能，因此2.0是100%向后兼容的。

也就是說(shuō)，代碼庫(kù)是一樣的，API也是一樣的，編寫模型的方式也是一樣的。

更多的技術(shù)支持TorchDynamo

使用Python框架評(píng)估鉤子安全地捕獲PyTorch程序，這是團(tuán)隊(duì)5年來(lái)在graph capture方面研發(fā)的一項(xiàng)重大創(chuàng)新。

AOTAutograd

重載了PyTorch的autograd引擎，作為一個(gè)追蹤的autodiff，用于生成超前的反向追蹤。

PrimTorch

將約2000多個(gè)PyTorch運(yùn)算符歸納為約250個(gè)原始運(yùn)算符的封閉集，開發(fā)人員可以針對(duì)這些運(yùn)算符構(gòu)建一個(gè)完整的PyTorch后端。大大降低了編寫PyTorch功能或后端的障礙。

TorchInductor

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

值得注意的是，TorchDynamo、AOTAutograd、PrimTorch和TorchInductor都是用Python編寫的，并支持動(dòng)態(tài)形狀。

更快的訓(xùn)練速度

通過(guò)引入新的編譯模式「torch.compile」，PyTorch 2.0用一行代碼，就可以加速模型的訓(xùn)練。

這里不用任何技巧，只需運(yùn)行torch.compile()即可，僅此而已：

opt_module = torch.compile(module)

為了驗(yàn)證這些技術(shù)，團(tuán)隊(duì)精心打造了測(cè)試基準(zhǔn)，包括圖像分類、物體檢測(cè)、圖像生成等任務(wù)，以及各種NLP任務(wù)，如語(yǔ)言建模、問(wèn)答、序列分類、推薦系統(tǒng)和強(qiáng)化學(xué)習(xí)。其中，這些基準(zhǔn)可以分為三類：

來(lái)自HuggingFace Transformers的46個(gè)模型來(lái)自TIMM的61個(gè)模型：Ross Wightman收集的最先進(jìn)的PyTorch圖像模型來(lái)自TorchBench的56個(gè)模型：github的一組流行代碼庫(kù)

測(cè)試結(jié)果表明，在這163個(gè)跨越視覺(jué)、NLP和其他領(lǐng)域的開源模型上，訓(xùn)練速度得到了38%-76%的提高。

在NVIDIA A100 GPU上的對(duì)比

此外，團(tuán)隊(duì)還在一些流行的開源PyTorch模型上進(jìn)行了基準(zhǔn)測(cè)試，并獲得了從30%到2倍的大幅加速。

開發(fā)者Sylvain Gugger表示：「只需添加一行代碼，PyTorch 2.0就能在訓(xùn)練Transformers模型時(shí)實(shí)現(xiàn)1.5倍到2.0倍的速度提升。這是自混合精度訓(xùn)練問(wèn)世以來(lái)最令人興奮的事情！」

技術(shù)概述

PyTorch的編譯器可以分解成三個(gè)部分：

圖的獲取圖的降低圖的編譯

其中，在構(gòu)建PyTorch編譯器時(shí)，圖的獲取是更難的挑戰(zhàn)。

TorchDynamo

今年年初，團(tuán)隊(duì)便開始了TorchDynamo的工作，這種方法使用了PEP-0523中引入的CPython功能，稱為框架評(píng)估API。

為此，團(tuán)隊(duì)采取了一種數(shù)據(jù)驅(qū)動(dòng)的方法來(lái)驗(yàn)證TorchDynamo在graph capture上的有效性——通過(guò)使用7000多個(gè)用PyTorch編寫的Github項(xiàng)目，來(lái)作為驗(yàn)證集。

結(jié)果顯示，TorchDynamo在99%的時(shí)間里都能正確、安全地進(jìn)行g(shù)raph capture，而且開銷可以忽略不計(jì)。

TorchInductor

對(duì)于PyTorch 2.0的新編譯器后端，團(tuán)隊(duì)從用戶如何編寫高性能的自定義內(nèi)核中得到了靈感：越來(lái)越多地使用Triton語(yǔ)言。

TorchInductor使用Pythonic定義的逐個(gè)循環(huán)級(jí)別的IR來(lái)自動(dòng)將PyTorch模型映射到GPU上生成的Triton代碼和CPU上的C++/OpenMP。

TorchInductor的核心循環(huán)級(jí)IR只包含大約50個(gè)運(yùn)算符，而且它是用Python實(shí)現(xiàn)的，這使得它很容易得到擴(kuò)展。

AOTAutograd

想要加速訓(xùn)練，就不僅需要捕獲用戶級(jí)代碼，而且還要捕獲反向傳播。

AOTAutograd可以利用PyTorch的torch_dispatch擴(kuò)展機(jī)制來(lái)追蹤Autograd引擎，「提前」捕獲反向傳播，進(jìn)而能夠使用TorchInductor來(lái)加速前向和后向通道。

PrimTorch

PyTorch有1200多個(gè)運(yùn)算符，如果考慮到每個(gè)運(yùn)算符的各種重載，則有2000多個(gè)。因此，編寫后端或跨領(lǐng)域的功能成為一項(xiàng)耗費(fèi)精力的工作。

在PrimTorch項(xiàng)目中，團(tuán)隊(duì)定義了兩個(gè)更小更穩(wěn)定的運(yùn)算符集：

Prim ops有大約~250個(gè)運(yùn)算符，適合于編譯器。由于足夠低級(jí)，因此只需將它們?nèi)诤显谝黄鹨垣@得良好的性能。ATen ops有大約~750個(gè)典型的運(yùn)算符，適合于按原樣輸出。這些適合于已經(jīng)在ATen級(jí)別上集成的后端，或者沒(méi)有編譯的后端，從而恢復(fù)像Prim ops這樣的低級(jí)別運(yùn)算符集的性能。動(dòng)態(tài)形狀

在研究支持PyTorch代碼通用性的必要條件時(shí)，一個(gè)關(guān)鍵要求是支持動(dòng)態(tài)形狀，并允許模型接受不同大小的張量，而不會(huì)在每次形狀變化時(shí)引起重新編譯。

在不支持動(dòng)態(tài)形狀的情況下，一個(gè)常見(jiàn)的解決方法是將其填充到最接近的2次方。然而，正如我們從下面的圖表中所看到的，它產(chǎn)生了大量的性能開銷，同時(shí)也帶來(lái)了明顯更長(zhǎng)的編譯時(shí)間。

現(xiàn)在，有了對(duì)動(dòng)態(tài)形狀的支持，PyTorch 2.0也就獲得了比Eager高出了最多40%的性能。

最后，在PyTorch 2.x的路線圖中，團(tuán)隊(duì)希望在性能和可擴(kuò)展性方面進(jìn)一步推動(dòng)編譯模式的發(fā)展。

關(guān)鍵詞： 進(jìn)一步推動(dòng) 可以利用 忽略不計(jì) 安全地進(jìn)行 數(shù)據(jù)驅(qū)動(dòng)