一行代碼，煉丹 2 倍速！PyTorch 2.0 驚喜問世，LeCun 激情轉發(fā)

12 月 2 日，PyTorch 2.0 正式發(fā)布！

這次的更新不僅將 PyTorch 的性能推到了新的高度，同時也加入了對動態(tài)形狀和分布式的支持。

此外，2.0 系列還會將 PyTorch 的部分代碼從 C++ 移回 Python。

目前，PyTorch 2.0 還處在測試階段，預計第一個穩(wěn)定版本會在 2023 年 3 月初面世。

PyTorch 2.x：更快、更 Python！

在過去的幾年里，PyTorch 從 1.0 到最近的 1.13 進行了創(chuàng)新和迭代，并轉移到新成立的 PyTorch 基金會，成為 Linux 基金會的一部分。

當前版本的 PyTorch 所面臨的挑戰(zhàn)是，eager-mode 難以跟上不斷增長的 GPU 帶寬和更瘋狂的模型架構。

而 PyTorch 2.0 的誕生，將從根本上改變和提升了 PyTorch 在編譯器級別下的運行方式。

眾所周知，PyTorch 中的（Py）來自于數(shù)據(jù)科學中廣泛使用的開源 Python 編程語言。

然而，PyTorch 的代碼卻并沒有完全采用 Python，而是把一部分交給了 C++。

不過，在今后的 2.x 系列中，PyTorch 項目團隊計劃將與 torch.nn 有關的代碼移回到 Python 中。

除此之外，由于 PyTorch 2.0 是一個完全附加的（和可選的）功能，因此 2.0 是 100% 向后兼容的。

也就是說，代碼庫是一樣的，API 也是一樣的，編寫模型的方式也是一樣的。

更多的技術支持

• TorchDynamo

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

• AOTAutograd

重載了 PyTorch 的 autograd 引擎，作為一個追蹤的 autodiff，用于生成超前的反向追蹤。

• PrimTorch

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

• TorchInductor

一個深度學習編譯器，可以為多個加速器和后端生成快速代碼。對于英偉達的 GPU，它使用 OpenAI Triton 作為關鍵構建模塊。

值得注意的是，TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 都是用 Python 編寫的，并支持動態(tài)形狀。

更快的訓練速度

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

這里不用任何技巧，只需運行 torch.compile () 即可，僅此而已：

opt_module = torch.compile(module)

為了驗證這些技術，團隊精心打造了測試基準，包括圖像分類、物體檢測、圖像生成等任務，以及各種 NLP 任務，如語言建模、問答、序列分類、推薦系統(tǒng)和強化學習。其中，這些基準可以分為三類：

• 來自 HuggingFace Transformers 的 46 個模型

• 來自 TIMM 的 61 個模型：Ross Wightman 收集的最先進的 PyTorch 圖像模型

• 來自 TorchBench 的 56 個模型：github 的一組流行代碼庫

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

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

開發(fā)者 Sylvain Gugger 表示：「只需添加一行代碼，PyTorch 2.0 就能在訓練 Transformers 模型時實現(xiàn) 1.5 倍到 2.0 倍的速度提升。這是自混合精度訓練問世以來最令人興奮的事情！」

技術概述

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

• 圖的獲取

• 圖的降低

• 圖的編譯

其中，在構建 PyTorch 編譯器時，圖的獲取是更難的挑戰(zhàn)。

TorchDynamo

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

為此，團隊采取了一種數(shù)據(jù)驅動的方法來驗證 TorchDynamo 在 graph capture 上的有效性 —— 通過使用 7000 多個用 PyTorch 編寫的 Github 項目，來作為驗證集。

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

TorchInductor

對于 PyTorch 2.0 的新編譯器后端，團隊從用戶如何編寫高性能的自定義內(nèi)核中得到了靈感：越來越多地使用 Triton 語言。

TorchInductor 使用 Pythonic 定義的逐個循環(huán)級別的 IR 來自動將 PyTorch 模型映射到 GPU 上生成的 Triton 代碼和 CPU 上的 C++/OpenMP。

TorchInductor 的核心循環(huán)級 IR 只包含大約 50 個運算符，而且它是用 Python 實現(xiàn)的，這使得它很容易得到擴展。

AOTAutograd

想要加速訓練，就不僅需要捕獲用戶級代碼，而且還要捕獲反向傳播。

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

PrimTorch

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

在 PrimTorch 項目中，團隊定義了兩個更小更穩(wěn)定的運算符集：

Prim ops 有大約~250 個運算符，適合于編譯器。由于足夠低級，因此只需將它們?nèi)诤显谝黄鹨垣@得良好的性能。

ATen ops 有大約~750 個典型的運算符，適合于按原樣輸出。這些適合于已經(jīng)在 ATen 級別上集成的后端，或者沒有編譯的后端，從而恢復像 Prim ops 這樣的低級別運算符集的性能。

動態(tài)形狀

在研究支持 PyTorch 代碼通用性的必要條件時，一個關鍵要求是支持動態(tài)形狀，并允許模型接受不同大小的張量，而不會在每次形狀變化時引起重新編譯。

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

現(xiàn)在，有了對動態(tài)形狀的支持，PyTorch 2.0 也就獲得了比 Eager 高出了最多 40% 的性能。

最后，在 PyTorch 2.x 的路線圖中，團隊希望在性能和可擴展性方面進一步推動編譯模式的發(fā)展。

參考資料：

• https://pytorch.org/get-started/pytorch-2.0/

本文來自微信公眾號：新智元 （ID：AI_era），編輯：好困

廣告聲明：文內(nèi)含有的對外跳轉鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。