【新智元導讀】PyTorch 2.0正式發布。

PyTorch 2.0正式版終于來了！

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

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

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

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

亮點總結

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

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

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

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

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

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

編譯，還是編譯！

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

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

TorchDynamo

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

AOTAutograd

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

PrimTorch

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

4. TorchInductor

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

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

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

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

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

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

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

NVIDIA A100 GPU eager mode torch.compile 針對不同模型的提速表現

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

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

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

訓練LLM提速2倍

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

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

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

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

Transformers是幫助實現現代生成式人工智能時代的基礎技術，包括GPT-3以及GPT-4這樣的OpenAI模型。

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

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

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

新版本能夠實現對訓練和推理的高性能支持，使用定制的內核架構來處理縮放點積注意力(SPDA) ，擴展了推理的快速路徑架構。

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

-看到速度明顯提升；

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

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

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

除了現有的Transformer API，開發者還可以通過調用新的scaled_dot_product_attention()操作符直接使用縮放點積注意力

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

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

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

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

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

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

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