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

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

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

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

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

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

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

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

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

代碼部分截圖

完整代碼地址：

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

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

部分結(jié)果截圖

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

使用 Trainer 類(lèi)?

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

DeepSpeed

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

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

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

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