一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

王林
發布: 2023-04-13 08:49:02
轉載
1453 人瀏覽過

12月2日,PyTorch 2.0正式發布!

這次的更新不僅將PyTorch的效能推到了新的高度,同時也加入了對動態形狀和分散式的支援。

此外,2.0系列也會將PyTorch的部分程式碼從C 移回Python。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

目前,PyTorch 2.0還在測試階段,預期第一個穩定版本會在2023年3月初面世。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

PyTorch 2.x:更快、更Python!

在過去的幾年裡,PyTorch從1.0到最近的1.13進行了創新和迭代,並轉移到新成立的PyTorch基金會,成為Linux基金會的一部分。

目前版本的PyTorch所面臨的挑戰是,eager-mode難以跟上不斷成長的GPU頻寬和更瘋狂的模型架構。

而PyTorch 2.0的誕生,將從根本上改變並提升了PyTorch在編譯器層級下的運作方式。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

眾所周知,PyTorch中的(Py)來自於資料科學中廣泛使用的開源Python程式語言。

然而,PyTorch的程式碼卻並沒有完全採用Python,而是把一部分交給了C 。

不過,在今後的2.x系列中,PyTorch專案團隊計畫將與torch.nn相關的程式碼移回Python中。

除此之外,由於PyTorch 2.0是一個完全附加的(和可選的)功能,因此2.0是100%向後相容的。

也就是說,程式碼庫是一樣的,API也是一樣的,寫模型的方式也是一樣的。

更多的技術支援

  • TorchDynamo

使用Python框架評估鉤子安全地捕獲PyTorch程序,這是團隊5年來在graph capture方面研發的一項重大創新。

  • AOTAutograd

重載了PyTorch的autograd引擎,作為一個追蹤的autodiff,用於產生超前的反向追蹤。

  • PrimTorch

將約2000多個PyTorch運算子歸納為約250個原始運算符的封閉集,開發人員可以針對這些運算子建構一個完整的PyTorch後端。大大降低了編寫PyTorch功能或後端的障礙。

  • TorchInductor

一個深度學習編譯器,可以為多個加速器和後端產生快速程式碼。對於英偉達的GPU,它使用OpenAI Triton作為關鍵構建模組。

值得注意的是,TorchDynamo、AOTAutograd、PrimTorch和TorchInductor都是用Python寫的,並且支援動態形狀。

更快的訓練速度

透過引入新的編譯模式「torch.compile」,PyTorch 2.0用一行程式碼,就可以加速模型的訓練。

這裡不用任何技巧,只需執行torch.compile()即可,僅此而已:

opt_module = torch.compile(module)
登入後複製

為了驗證這些技術,團隊精心打造了測試基準,包括影像分類、物件偵測、影像生成等任務,以及各種NLP任務,如語言建模、問答、序列分類、推薦系統和強化學習。其中,這些基準可以分為三類:

  • 來自HuggingFace Transformers的46個模型
  • 來自TIMM的61個模型:Ross Wightman收集的最先進的PyTorch圖像模型
  • 來自TorchBench的56個模型:github的一組流行程式碼庫

測試結果表明,在這163個跨越視覺、NLP和其他領域的在開源模型上,訓練速度得到了38%-76%的提升。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

在NVIDIA A100 GPU上的比較

此外,團隊還在一些流行的開源PyTorch模型上進行了基準測試,並獲得了從30%到2倍的大幅加速。

開發者Sylvain Gugger表示:「只要加入一行程式碼,PyTorch 2.0就能在訓練Transformers模型時實現1.5倍到2.0倍的速度提升。這是自混合精度訓練問世以來最令人興奮的事情!」

技術概述

PyTorch的編譯器可以分解成三個部分:

  • 圖的取得
  • 圖的降低
  • 圖的編譯

#其中,在建構PyTorch編譯器時,圖的取得是更難的挑戰。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

TorchDynamo

今年年初,團隊便開始了TorchDynamo的工作,這種方法使用了PEP-0523中引入的CPython功能,稱為框架評估API。

為此,團隊採取了一種數據驅動的方法來驗證TorchDynamo在graph capture上的有效性——透過使用7000多個用PyTorch編寫的Github項目,來作為驗證集。

結果顯示,TorchDynamo在99%的時間裡都能正確、安全地進行graph capture,而且開銷可以忽略不計。

TorchInductor

對於PyTorch 2.0的新編譯器後端,團隊從使用者如何編寫高效能的自訂核心中得到了靈感:越來越多地使用Triton語言。

TorchInductor使用Pythonic定義的逐個循環層級的IR來自動將PyTorch模型對應到GPU上產生的Triton程式碼和CPU上的C /OpenMP。

TorchInductor的核心循環級IR只包含大約50個運算符,而且它是用Python實現的,這使得它很容易擴展。

AOTAutograd

想要加速訓練,不僅需要捕捉使用者級程式碼,還要捕捉反向傳播。

AOTAutograd可以利用PyTorch的torch_dispatch擴展機制來追蹤Autograd引擎,「提前」捕捉反向傳播,進而能夠使用TorchInductor來加速前向和後向通道。

PrimTorch

PyTorch有1200多個運算符,如果考慮到每個運算符的各種重載,則有2000多個個。因此,編寫後端或跨領域的功能成為一項耗費精力的工作。

在PrimTorch專案中,團隊定義了兩個更小、更穩定的運算子集:

#
  • Prim ops有大約~250個運算符,適合編譯器。由於足夠低級,因此只需將它們融合在一起以獲得良好的性能。
  • ATen ops有大約~750個典型的運算符,適合於原樣輸出。這些適合已經在ATen層級上整合的後端,或是沒有編譯的後端,從而恢復像Prim ops這樣的低階運算子集的效能。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

動態形狀

#在研究支援PyTorch程式碼通用性的必要條件時,一個關鍵要求是支援動態形狀,並允許模型接受不同大小的張量,而不會在每次形狀變化時引起重新編譯。

在不支援動態形狀的情況下,一個常見的解決方法是將其填入最接近的2次方。然而,正如我們從下面的圖表中所看到的,它產生了大量的效能開銷,同時也帶來了明顯更長的編譯時間。

現在,有了對動態形狀的支持,PyTorch 2.0也就獲得了比Eager高出了最多40%的性能。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

最後,在PyTorch 2.x的路線圖中,團隊希望在效能和可擴展性方面進一步推動編譯模式的發展。

一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發

#

以上是一行程式碼,煉丹2倍速! PyTorch 2.0驚喜問世,LeCun激情轉發的詳細內容。更多資訊請關注PHP中文網其他相關文章!

相關標籤:
來源:51cto.com
本網站聲明
本文內容由網友自願投稿,版權歸原作者所有。本站不承擔相應的法律責任。如發現涉嫌抄襲或侵權的內容,請聯絡admin@php.cn
最新問題
熱門教學
更多>
最新下載
更多>
網站特效
網站源碼
網站素材
前端模板