3 月 19 日消息，PyTorch 2.0 穩(wěn)定版現(xiàn)已發(fā)布。跟先前 1.0 版本相比，2.0 有了顛覆式的變化。在 PyTorch 2.0 中，最大的改進主要是 API 的 torch.compile，新編譯器比先前「eager mode」所提供的即時生成代碼的速度快得多，性能得以進一步提升。

新版本更新之處包括穩(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 改進。


(資料圖片)

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

除了 PyTorch2.0，研發(fā)團隊還發(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í)編譯器來生成性能代碼并隱藏低級硬件細節(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 多個運算符。

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

其他一些跨 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ā)團隊表示他們希望整個行業(yè)都能負擔(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) 對模型進行預(yù)處理。

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

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

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

這個數(shù)據(jù)來自 PyTorch 基金會在 Nvidia A100 GPU 上使用 PyTorch 2.0 對 163 個開源模型進行的基準測試，其中包括圖像分類、目標檢測、圖像生成等任務(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 默認后端 TorchInductor 已經(jīng)支持 CPU 和 NVIDIA Volta 和 Ampere GP，暫不支持其他 GPU、XPU 或舊的 NVIDIA GPU。

標簽：