環(huán)球熱門:改變幾行代碼，PyTorch 煉丹速度狂飆、模型優(yōu)化時間大減

選自 Sebastian Raschka 博客

【資料圖】

機(jī)器之心編譯

機(jī)器之心編輯部

關(guān)于 PyTorch 煉丹，本文作者表示：「如果你有 8 個 GPU，整個訓(xùn)練過程只需要 2 分鐘，實(shí)現(xiàn) 11.5 倍的性能加速。」

如何提升 PyTorch「煉丹」速度？

最近，知名機(jī)器學(xué)習(xí)與 AI 研究者 Sebastian Raschka 向我們展示了他的絕招。據(jù)他表示，他的方法在不影響模型準(zhǔn)確率的情況下，僅僅通過改變幾行代碼，將 BERT 優(yōu)化時間從 22.63 分鐘縮減到 3.15 分鐘，訓(xùn)練速度足足提升了 7 倍。

作者更是表示，如果你有 8 個 GPU 可用，整個訓(xùn)練過程只需要 2 分鐘，實(shí)現(xiàn) 11.5 倍的性能加速。

下面我們來看看他到底是如何實(shí)現(xiàn)的。

讓 PyTorch 模型訓(xùn)練更快

首先是模型，作者采用 DistilBERT 模型進(jìn)行研究，它是 BERT 的精簡版，與 BERT 相比規(guī)模縮小了 40%，但性能幾乎沒有損失。其次是數(shù)據(jù)集，訓(xùn)練數(shù)據(jù)集為大型電影評論數(shù)據(jù)集 IMDB Large Movie Review，該數(shù)據(jù)集總共包含 50000 條電影評論。作者將使用下圖中的 c 方法來預(yù)測數(shù)據(jù)集中的影評情緒。

基本任務(wù)交代清楚后，下面就是 PyTorch 的訓(xùn)練過程。為了讓大家更好地理解這項(xiàng)任務(wù)，作者還貼心地介紹了一下熱身練習(xí)，即如何在 IMDB 電影評論數(shù)據(jù)集上訓(xùn)練 DistilBERT 模型。如果你想自己運(yùn)行代碼，可以使用相關(guān)的 Python 庫設(shè)置一個虛擬環(huán)境，如下所示：

相關(guān)軟件的版本如下：

現(xiàn)在省略掉枯燥的數(shù)據(jù)加載介紹，只需要了解本文將數(shù)據(jù)集劃分為 35000 個訓(xùn)練示例、5000 個驗(yàn)證示例和 10000 個測試示例。需要的代碼如下：

代碼部分截圖

完整代碼地址：

https://github.com/rasbt/faster-pytorch-blog/blob/main/1_pytorch-distilbert.py

然后在 A100 GPU 上運(yùn)行代碼，得到如下結(jié)果：

部分結(jié)果截圖

正如上述代碼所示，模型從第 2 輪到第 3 輪開始有一點(diǎn)過擬合，驗(yàn)證準(zhǔn)確率從 92.89% 下降到了 92.09%。在模型運(yùn)行了 22.63 分鐘后進(jìn)行微調(diào)，最終的測試準(zhǔn)確率為 91.43%。

使用 Trainer 類

接下來是改進(jìn)上述代碼，改進(jìn)部分主要是把 PyTorch 模型包裝在 LightningModule 中，這樣就可以使用來自 Lightning 的 Trainer 類。部分代碼截圖如下：

完整代碼地址：https://github.com/rasbt/faster-pytorch-blog/blob/main/2_pytorch-with-trainer.py

上述代碼建立了一個 LightningModule，它定義了如何執(zhí)行訓(xùn)練、驗(yàn)證和測試。相比于前面給出的代碼，主要變化是在第 5 部分（即 ### 5 Finetuning），即微調(diào)模型。與以前不同的是，微調(diào)部分在 LightningModel 類中包裝了 PyTorch 模型，并使用 Trainer 類來擬合模型。

之前的代碼顯示驗(yàn)證準(zhǔn)確率從第 2 輪到第 3 輪有所下降，但改進(jìn)后的代碼使用了 ModelCheckpoint 以加載最佳模型。在同一臺機(jī)器上，這個模型在 23.09 分鐘內(nèi)達(dá)到了 92% 的測試準(zhǔn)確率。

需要注意，如果禁用 checkpointing 并允許 PyTorch 以非確定性模式運(yùn)行，本次運(yùn)行最終將獲得與普通 PyTorch 相同的運(yùn)行時間（時間為 22.63 分而不是 23.09 分）。

自動混合精度訓(xùn)練

進(jìn)一步，如果 GPU 支持混合精度訓(xùn)練，可以開啟 GPU 以提高計(jì)算效率。作者使用自動混合精度訓(xùn)練，在 32 位和 16 位浮點(diǎn)之間切換而不會犧牲準(zhǔn)確率。

在這一優(yōu)化下，使用 Trainer 類，即能通過一行代碼實(shí)現(xiàn)自動混合精度訓(xùn)練：

上述操作可以將訓(xùn)練時間從 23.09 分鐘縮短到 8.75 分鐘，這幾乎快了 3 倍。測試集的準(zhǔn)確率為 92.2%，甚至比之前的 92.0% 還略有提高。

使用 Torch.Compile 靜態(tài)圖

最近 PyTorch 2.0 公告顯示，PyTorch 團(tuán)隊(duì)引入了新的 toch.compile 函數(shù)。該函數(shù)可以通過生成優(yōu)化的靜態(tài)圖來加速 PyTorch 代碼執(zhí)行，而不是使用動態(tài)圖運(yùn)行 PyTorch 代碼。

由于 PyTorch 2.0 尚未正式發(fā)布，因而必須先要安裝 torchtriton，并更新到 PyTorch 最新版本才能使用此功能。

然后通過添加這一行對代碼進(jìn)行修改：

在 4 塊 GPU 上進(jìn)行分布式數(shù)據(jù)并行

上文介紹了在單 GPU 上加速代碼的混合精度訓(xùn)練，接下來介紹多 GPU 訓(xùn)練策略。下圖總結(jié)了幾種不同的多 GPU 訓(xùn)練技術(shù)。

想要實(shí)現(xiàn)分布式數(shù)據(jù)并行，可以通過 DistributedDataParallel 來實(shí)現(xiàn)，只需修改一行代碼就能使用 Trainer。

經(jīng)過這一步優(yōu)化，在 4 個 A100 GPU 上，這段代碼運(yùn)行了 3.52 分鐘就達(dá)到了 93.1% 的測試準(zhǔn)確率。

DeepSpeed

最后，作者探索了在 Trainer 中使用深度學(xué)習(xí)優(yōu)化庫 DeepSpeed 以及多 GPU 策略的結(jié)果。首先必須安裝 DeepSpeed 庫：

接著只需更改一行代碼即可啟用該庫：

這一波下來，用時 3.15 分鐘就達(dá)到了 92.6% 的測試準(zhǔn)確率。不過 PyTorch 也有 DeepSpeed 的替代方案：fully-sharded DataParallel，通過 strategy="fsdp" 調(diào)用，最后花費(fèi) 3.62 分鐘完成。

以上就是作者提高 PyTorch 模型訓(xùn)練速度的方法，感興趣的小伙伴可以跟著原博客嘗試一下，相信你會得到想要的結(jié)果。

原文鏈接：https://sebastianraschka.com/blog/2023/pytorch-faster.html

探尋隱私計(jì)算最新行業(yè)技術(shù)，「首屆隱語開源社區(qū)開放日」報名啟程

春暖花開之際，誠邀廣大技術(shù)開發(fā)者 & 產(chǎn)業(yè)用戶相聚活動現(xiàn)場，體驗(yàn)數(shù)智時代的隱私計(jì)算生態(tài)建設(shè)之旅，一站構(gòu)建隱私計(jì)算產(chǎn)業(yè)體系知識：

隱私計(jì)算領(lǐng)域焦點(diǎn)之性

分布式計(jì)算系統(tǒng)的短板與升級策略

隱私計(jì)算跨平臺互聯(lián)互通

隱語開源框架金融行業(yè)實(shí)戰(zhàn)經(jīng)驗(yàn)

3 月 29 日，北京 · 798 機(jī)遇空間，隱語開源社區(qū)開放日，期待線下面基。