【導讀】PyTorch 2.0正式發(fā)布。

PyTorch 2.0正式版終于來了！

去年12月，PyTorch基金會在PyTorch Conference 2022上發(fā)布了PyTorch 2.0的第一個預覽版本。

跟先前1.0版本相比，2.0有了顛覆式的變化。在PyTorch 2.0中，最大的改進是torch.compile。

新的編譯器比以前PyTorch 1.0中默認的「eager mode」所提供的即時生成代碼的速度快得多，讓PyTorch性能進一步提升。

除了2.0之外，還發(fā)布了一系列PyTorch域庫的beta更新，包括那些在樹中的庫，以及包括 TorchAudio、TorchVision和TorchText在內(nèi)的獨立庫。TorchX的更新也同時發(fā)布，可以提供社區(qū)支持模式。

亮點總結(jié)

-torch.compile是PyTorch 2.0的主要API，它包裝并返回編譯后的模型，torch.compile是一個完全附加（和可選）的特性，因此2.0版本是100%向后兼容的。

-作為torch.compile的基礎(chǔ)技術(shù)，帶有Nvidia和AMD GPU的TorchInductor將依賴OpenAI Triton深度學習編譯器來生成高性能代碼，并隱藏低級硬件細節(jié)。OpenAI Triton生成的內(nèi)核實現(xiàn)的性能，與手寫內(nèi)核和cublas等專門的cuda庫相當。

-Accelerated Transformers引入了對訓練和推理的高性能支持，使用自定義內(nèi)核架構(gòu)實現(xiàn)縮放點積注意力 （SPDA）。API與torch.compile （） 集成，模型開發(fā)人員也可以通過調(diào)用新的scaled_dot_product_attention （） 運算符，直接使用縮放的點積注意力內(nèi)核。

-Metal Performance Shaders （MPS） 后端在Mac平臺上提供GPU加速的PyTorch訓練，并增加了對前60個最常用操作的支持，覆蓋了300多個操作符。

-Amazon AWS優(yōu)化了基于AWS Graviton3的C7g實例上的PyTorch CPU推理。與之前的版本相比，PyTorch 2.0提高了Graviton的推理性能，包括對Resnet50和Bert的改進。

-跨TensorParallel、DTensor、2D parallel、TorchDynamo、AOTAutograd、PrimTorch和TorchInductor的新原型功能和技術(shù)。

編譯，還是編譯！

PyTorch 2.0的最新編譯器技術(shù)包括：TorchDynamo、AOTAutograd、PrimTorch和TorchInductor。所有這些都是用Python開發(fā)的，而不是C++（Python與之兼容）。

并且還支持dynamic shape，無需重新編譯就能發(fā)送不同大小的向量，靈活且易學。

TorchDynamo

它可以借助Python Frame Evaluation Hooks，安全地獲取PyTorch程序，這項重大創(chuàng)新是PyTorch過去 5 年來在安全圖結(jié)構(gòu)捕獲 （safe graph capture） 方面的研發(fā)成果匯總。

AOTAutograd

重載PyTorch autograd engine，作為一個 tracing autodiff，用于生成超前的backward trace。

PrimTorch

將 2000+ PyTorch 算子歸納為約 250 個 primitive operator 閉集 （closed set），開發(fā)者可以針對這些算子構(gòu)建一個完整的 PyTorch 后端。PrimTorch 大大簡化了編寫 PyTorch 功能或后端的流程。

4. TorchInductor

TorchInductor一個深度學習編譯器，可以為多個加速器和后端生成 fast code。對于 NVIDIA GPU，它使用 OpenAI Triton 作為關(guān)鍵構(gòu)建模塊。

PyTorch基金會稱，2.0的推出會推動「從C++回到Python」，并補充說這是PyTorch的一個實質(zhì)性的新方向。

「從第一天起，我們就知道「eager execution」的性能限制。2017年7月，我們開始了第一個研究項目，為PyTorch開發(fā)一個編譯器。編譯器需要使PyTorch程序快速運行，但不能以PyTorch的體驗為代價，還要保留靈活易用性，這樣的話可以支持研究人員在不同探索階段使用動態(tài)的模型和程序?！?/p>

當然了，非編譯的「eager mode」使用動態(tài)即時代碼生成器，在2.0中仍然可用。開發(fā)者可以使用porch.compile命令迅速升級到編譯模式，只需要增加一行代碼。

用戶可以看到2.0的編譯時間比1.0提高43%。

這個數(shù)據(jù)來自PyTorch基金會在Nvidia A100 GPU上使用PyTorch 2.0對163個開源模型進行的基準測試，其中包括包括圖像分類、目標檢測、圖像生成等任務，以及各種 NLP 任務。

這些Benchmark分為三類：HuggingFace Tranformers、TIMM和TorchBench。

NVIDIA A100 GPU eager mode torch.compile 針對不同模型的提速表現(xiàn)

據(jù)PyTorch基金會稱，新編譯器在使用Float32精度模式時運行速度提高了21%，在使用自動混合精度（AMP）模式時運行速度提高了51%。

在這163個模型中，torch.compile可以在93%模型上正常運行。

「在PyTorch 2.x的路線圖中，我們希望在性能和可擴展性方面讓編譯模式越走越遠。有一些工作還沒有開始。有些工作因為帶寬不夠而辦法落地。」

訓練LLM提速2倍

此外，性能是PyTorch 2.0的另一個主要重點，也是開發(fā)人員一直不吝于宣傳的一個重點。

事實上，新功能的亮點之一是Accelerated Transformers，之前被稱為Better Transformers。

另外，PyTorch 2.0正式版包含了一個新的高性能PyTorch TransformAPI實現(xiàn)。

PyTorch項目的一個目標，是讓最先進的transformer模型的訓練和部署更加容易、快速。

Transformers是幫助實現(xiàn)現(xiàn)代生成式人工智能時代的基礎(chǔ)技術(shù)，包括GPT-3以及GPT-4這樣的OpenAI模型。

在PyTorch 2.0 Accelerated Transformers中，使用了自定義內(nèi)核架構(gòu)的方法（也被稱為縮放點積注意力SDPA），為訓練和推理提供高性能的支持。

由于有多種類型的硬件可以支持Transformers，PyTorch 2.0可以支持多個SDPA定制內(nèi)核。更進一步，PyTorch集成了自定義內(nèi)核選擇邏輯，將為給定的模型和硬件類型挑選最高性能的內(nèi)核。

加速的影響非同小可，因為它有助于使開發(fā)人員比以前的PyTorch迭代更快地訓練模型。

新版本能夠?qū)崿F(xiàn)對訓練和推理的高性能支持，使用定制的內(nèi)核架構(gòu)來處理縮放點積注意力（SPDA） ，擴展了推理的快速路徑架構(gòu)。

與fastpath架構(gòu)類似，定制內(nèi)核完全集成到PyTorch TransformerAPI中--因此，使用本地Transformer和MultiHeadAttention API將使用戶能夠：

-看到速度明顯提升；

-支持更多的用例，包括使用交叉注意的模型、Transformer解碼器和訓練模型；

-繼續(xù)將快速路徑推理用于固定和可變序列長度的變形器編碼器和自注意力機制的用例。

為了充分利用不同的硬件模型和Transformer用例，支持多個SDPA自定義內(nèi)核，自定義內(nèi)核選擇邏輯將為特定模型和硬件類型挑選最高性能的內(nèi)核。

除了現(xiàn)有的Transformer API，開發(fā)者還可以通過調(diào)用新的scaled_dot_product_attention（）操作符直接使用縮放點積注意力關(guān)注內(nèi)核，加速PyTorch 2 Transformers與torch.compile（）集成。

為了在使用模型的同時，還能獲得PT2編譯的額外加速（用于推理或訓練），可以使用model = torch.compile（model）對模型進行預處理。

目前，已經(jīng)使用自定義內(nèi)核和torch.compile（）的組合，在訓練Transformer模型，特別是使用加速的PyTorch 2 Transformer的大型語言模型方面取得實質(zhì)性加速提升。

使用自定義內(nèi)核和 torch.compile來提供對大型語言模型訓練顯著加速

HuggingFace Transformers的主要維護者Sylvain Gugger在PyTorch項目發(fā)表的一份聲明中寫道「只需添加一行代碼，PyTorch 2.0就能在訓練Transformers模型時提供1.5倍至2.0倍的速度。這是自混合精度訓練推出以來最令人興奮的事情！」

PyTorch和谷歌的TensorFlow是兩個最流行的深度學習框架。世界上有數(shù)千家機構(gòu)正在使用PyTorch開發(fā)深度學習應用程序，而且它的使用量正在不斷增加。

PyTorch 2.0的推出將有助于加速深度學習和人工智能應用的發(fā)展，Lightning AI的首席技術(shù)官和PyTorch Lightning的主要維護者之一Luca Antiga表示：

「PyTorch 2.0 體現(xiàn)了深度學習框架的未來。不需要用戶干預即可捕獲PyTorch 程序，開箱即用的程序生成，以及巨大的設(shè)備加速，這種可能性為人工智能開發(fā)人員打開了一個全新的維度。」

審核編輯 ：李倩