一行代碼，煉丹2倍速！PyTorch 2.0驚喜問世，LeCun激情轉發
12月2日，PyTorch 2.0正式發布！
這次的更新不僅將PyTorch的性能推到了新的高度，同時也加入了對動態形狀和分布式的支持。
此外，2.0系列還會將PyTorch的部分代碼從C++移回Python。
目前，PyTorch 2.0還處在測試階段，預計第一個穩定版本會在2023年3月初面世。
PyTorch 2.x：更快、更Python！
在過去的幾年里，PyTorch從1.0到最近的1.13進行了創新和迭代，并轉移到新成立的PyTorch基金會，成為Linux基金會的一部分。
當前版本的PyTorch所面臨的挑戰是，eager-mode難以跟上不斷增長的GPU帶寬和更瘋狂的模型架構。
而PyTorch 2.0的誕生，將從根本上改變和提升了PyTorch在編譯器級別下的運行方式。
眾所周知，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()即可，僅此而已：
為了驗證這些技術，團隊精心打造了測試基準，包括圖像分類、物體檢測、圖像生成等任務，以及各種NLP任務，如語言建模、問答、序列分類、推薦系統和強化學習。其中，這些基準可以分為三類：
- 來自HuggingFace Transformers的46個模型
- 來自TIMM的61個模型：Ross Wightman收集的最先進的PyTorch圖像模型
- 來自TorchBench的56個模型：github的一組流行代碼庫
測試結果表明，在這163個跨越視覺、NLP和其他領域的開源模型上，訓練速度得到了38%-76%的提高。
在NVIDIA A100 GPU上的對比
此外，團隊還在一些流行的開源PyTorch模型上進行了基準測試，并獲得了從30%到2倍的大幅加速。
開發者Sylvain Gugger表示：「只需添加一行代碼，PyTorch 2.0就能在訓練Transformers模型時實現1.5倍到2.0倍的速度提升。這是自混合精度訓練問世以來最令人興奮的事情！」
技術概述
PyTorch的編譯器可以分解成三個部分：
- 圖的獲取
- 圖的降低
- 圖的編譯
其中，在構建PyTorch編譯器時，圖的獲取是更難的挑戰。
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這樣的低級別運算符集的性能。
動態形狀
在研究支持PyTorch代碼通用性的必要條件時，一個關鍵要求是支持動態形狀，并允許模型接受不同大小的張量，而不會在每次形狀變化時引起重新編譯。
在不支持動態形狀的情況下，一個常見的解決方法是將其填充到最接近的2次方。然而，正如我們從下面的圖表中所看到的，它產生了大量的性能開銷，同時也帶來了明顯更長的編譯時間。
現在，有了對動態形狀的支持，PyTorch 2.0也就獲得了比Eager高出了最多40%的性能。
最后，在PyTorch 2.x的路線圖中，團隊希望在性能和可擴展性方面進一步推動編譯模式的發展。
