PyTorch 2.0穩(wěn)定版發(fā)布：性能進(jìn)一步提升

　　PyTorch 2.0穩(wěn)定版已經(jīng)正式發(fā)布了，和第一代的版本相比有了非常大的改進(jìn)，對API的torch.compile進(jìn)行了更好的優(yōu)化，而且在進(jìn)行代碼編輯的時候，PyTorch 2.0穩(wěn)定版也擁有更加快速的生成速度，讓你體驗到即時生成的快感。程序整體的性能也有非常大的提升。一起來看看關(guān)于PyTorch 2.0穩(wěn)定版的詳細(xì)介紹吧。

　　新版本更新之處包括穩(wěn)定版的 Accelerated Transformers（以前稱為 Better Transformers）；Beta 版包括作為 PyTorch 2.0 主要 API 的 torch.compile、作為 torch.nn.functional 一部分的 scaled_dot_product_attention 函數(shù)、MPS 后端、torch.func 模塊中的 functorch API；以及其他跨越各種推理、性能和訓(xùn)練優(yōu)化功能的 GPU 和 CPU 的 Beta / Prototype 改進(jìn)。

　　關(guān)于 torch.compile 的全面介紹和技術(shù)概況請見 2.0 入門頁面。

　　除了 PyTorch 2.0，研發(fā)團(tuán)隊還發(fā)布了 PyTorch 域庫的一系列 Beta 更新，包括 in-tree 的庫和 TorchAudio、TorchVision、TorchText 等獨立庫。此外，TorchX 轉(zhuǎn)向社區(qū)支持模式。

　　概括：

　　torch.compile 是 PyTorch 2.0 的主要 API，它能包裝并返回編譯后的模型。這個是一個完全附加（和可選）的功能，因此 PyTorch 2.0 根據(jù)定義是 100% 向后兼容的。

　　作為 torch.compile 的基礎(chǔ)技術(shù)，TorchInductor 與 Nvidia / AMD GPU 將依賴于 OpenAI Triton 深度學(xué)習(xí)編譯器來生成性能代碼并隱藏低級硬件細(xì)節(jié)。OpenAI triton 生成的內(nèi)核則實現(xiàn)了與手寫內(nèi)核和專用 cuda 庫 （如 cublas） 相當(dāng)?shù)男阅堋?/p>

　　Accelerated Transformers 引入了對訓(xùn)練和推理的高性能支持，使用自定義內(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 訓(xùn)練，并增加了對前 60 個最常用運算符的支持，覆蓋 300 多個運算符。

　　Amazon AWS 優(yōu)化了 AWS Graviton3 上的 PyTorch CPU 推理。與之前的版本相比，PyTorch 2.0 提高了 Graviton 的推理性能，包括針對 ResNet-50 和 BERT 的改進(jìn)。

　　其他一些跨 TensorParallel、DTensor、2D parallel、TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 的新 prototype 功能和方法。

　　要查看公開的 2.0、1.13 和 1.12 功能完整列表，請點擊此處。

　　穩(wěn)定功能

　　PyTorch 2.0 版本包括 PyTorch Transformer API 新的高性能實現(xiàn)，以前稱為「Better Transformer API」，現(xiàn)在更名為 「Accelerated PyTorch 2 Transformers」。

　　研發(fā)團(tuán)隊表示他們希望整個行業(yè)都能負(fù)擔(dān)得起訓(xùn)練和部署 SOTA Transformer 模型的成本。新版本引入了對訓(xùn)練和推理的高性能支持，使用自定義內(nèi)核架構(gòu)實現(xiàn)縮放點積注意力 （SPDA）。

　　與「快速路徑（fastpath）」架構(gòu)類似，自定義內(nèi)核完全集成到 PyTorch Transformer API 中 —— 因此，使用 Transformer 和 MultiHeadAttention API 將使用戶能夠：

　　明顯地看到顯著的速度提升；

　　支持更多用例，包括使用交叉注意力模型、Transformer 解碼器，并且可以用于訓(xùn)練模型；

　　繼續(xù)對固定和可變的序列長度 Transformer 編碼器和自注意力用例使用 fastpath 推理。

　　為了充分利用不同的硬件模型和 Transformer 用例，PyTorch 2.0 支持多個 SDPA 自定義內(nèi)核，自定義內(nèi)核選擇邏輯是為給定模型和硬件類型選擇最高性能的內(nèi)核。除了現(xiàn)有的 Transformer API 之外，模型開發(fā)人員還可以通過調(diào)用新的 scaled_dot_product_attention 運算來直接使用縮放點積注意力內(nèi)核。

　　要使用您的模型，同時受益于 pt2 編譯的額外加速 （用于推斷或訓(xùn)練），請使用 model = torch.compile （model） 對模型進(jìn)行預(yù)處理。

　　我們通過使用自定義內(nèi)核和 torch.compile 的組合，使用 Accelerated PyTorch 2 transformer 實現(xiàn)了訓(xùn)練 transformer 模型的大幅加速，特別是大語言模型。

　　將縮放點積注意力與自定義內(nèi)核和 torch.compile 結(jié)合使用可為訓(xùn)練大型語言模型（上圖以 nanoGPT 為例）提供顯著加速。

　　從官方數(shù)據(jù)可以看到，PyTorch 2.0 的編譯效率比 1.0 實現(xiàn)了大幅提高。

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

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

　　據(jù) PyTorch 基金會稱，新編譯器在使用 Float32 精度模式時運行速度提高了 21%，在使用自動混合精度（AMP）模式時運行速度提高了 51%。在這 163 個模型中，torch.compile 可以在 93% 模型上正常運行。

　　值得一提的是，官方在桌面級 GPU（如 NVIDIA 3090）上測量到的加速能力低于服務(wù)器級 GPU（如 A100）。到目前為止，PyTorch 2.0 默認(rèn)后端 TorchInductor 已經(jīng)支持 CPU 和 NVIDIA Volta 和 Ampere GP，暫不支持其他 GPU、XPU 或舊的 NVIDIA GPU。