機(jī)器之心報(bào)道
作者：杜偉、陳萍
繼1.0.0 版本推出不到兩個(gè)月的時(shí)間，grid.ai CEO、紐約大學(xué)博士 William Falcon 創(chuàng)建的 PyTorch Lightning 于近日宣布推出 1.1 版本。新版本新增了 sharded training 功能，在多 GPU 上訓(xùn)練深度學(xué)習(xí)（DL）模型時(shí)可以節(jié)省 50% 以上的內(nèi)存，并且沒(méi)有性能損失，也不需要更改代碼。
與 Facebook Research 的 FairScale 團(tuán)隊(duì)一道，PyTorch Lightning 團(tuán)隊(duì)在 1.1 版本中推出了 Sharded Training beta 版。在下面的博客文章中，PyTorch Lightning 團(tuán)隊(duì)的研究工程師 Sean Narenthiran 展示了只需要在 Lightning 訓(xùn)練器中增加單一的 flag，則在多 GPU 上訓(xùn)練 DL 模型時(shí)就會(huì)實(shí)現(xiàn)內(nèi)存的大幅度降低。
此外，作者還介紹了如何使用 NeMo 來(lái)預(yù)訓(xùn)練 Transformer LM，并實(shí)現(xiàn) 55% 的內(nèi)存提升，以及訓(xùn)練其他 PyTorch Lightning 賦能模型時(shí)進(jìn)一步的內(nèi)存降低。除了給出使用 NeMO Transformer LM 時(shí) NLP 中的結(jié)果之外，作者還分別展示了語(yǔ)音識(shí)別中使用 DeepSpeech 2 以及計(jì)算機(jī)視覺(jué)中訓(xùn)練 SwAV ResNet 和 iGPT 的結(jié)果。
PyTorch Lightning 團(tuán)隊(duì)正努力增添新的模型并行化技術(shù)并保證魯棒性，并且與 FairScale 團(tuán)隊(duì)展開(kāi)合作提升所有 PyTorch Lightning 研究中的模型擴(kuò)展性能。
更多使用技巧參考：https://pytorch-lightning.readthedocs.io/en/stable/multi_gpu.html#multi-gpu-training
更大的模型，更高的準(zhǔn)確率
語(yǔ)言建模趨向于更大的預(yù)訓(xùn)練模型，這種模型在下游任務(wù)中表現(xiàn)得更好。OpenAI 的 GPT-3 就是一個(gè)很好的例子，該模型有 1750 億個(gè)參數(shù)，在訓(xùn)練時(shí)需要大量的計(jì)算與優(yōu)化技巧。
比較了語(yǔ)言模型參數(shù)隨時(shí)間變化的曲線，GPT-3 繼續(xù)在規(guī)模上超越。（圖源：Microsoft）
訓(xùn)練大型模型時(shí)，內(nèi)存很寶貴。當(dāng)擴(kuò)展模型大小時(shí)，GPU 的內(nèi)存會(huì)被耗盡，而這又限制了訓(xùn)練模型的大小。這使得團(tuán)隊(duì)不斷嘗試更智能的內(nèi)存管理技術(shù)。
Lightning 的 Sharded Training
傳統(tǒng)分布式訓(xùn)練 VS Sharded Training。參數(shù)（P）在 GPU 之間拆分，以減少每個(gè) GPU 的內(nèi)存開(kāi)銷。Sharded Training 則拆分了優(yōu)化器狀態(tài)和梯度。
受微軟 Zero Redundancy Optimizer (ZeRO) 的啟發(fā)，Sharded Training 可以減少在多個(gè) GPU 上訓(xùn)練大型模型所需的內(nèi)存，訓(xùn)練過(guò)程中在 GPU 之間「切分」模型。Sharding 包括將參數(shù)分割到不同的設(shè)備上，減少每個(gè)設(shè)備所需的內(nèi)存。特別地，優(yōu)化器狀態(tài)和梯度可以獨(dú)立于模型進(jìn)行切分，并且可以減少所有架構(gòu)所需的內(nèi)存。
Sharded Training 是在 FairScale 基礎(chǔ)上構(gòu)建的，與 PyTorch 兼容并得到優(yōu)化。FairScale 是一個(gè) PyTorch 擴(kuò)展庫(kù)，用于高性能以及大規(guī)模訓(xùn)練模型和數(shù)據(jù)并行。除了切分技術(shù)之外，它具有層間和層內(nèi)并行性以及跨多個(gè) GPU 和主機(jī)拆分模型。
通過(guò)在 GPU 上進(jìn)行智能梯度和優(yōu)化器狀態(tài) sharding，可以分別將內(nèi)存成本（基于微軟論文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》的數(shù)據(jù)）降低大約 4 倍和 8 倍。這有利于所有模型，在所有模型架構(gòu)以及訓(xùn)練過(guò)程中提供較低的內(nèi)存使用率。需要注意的是，由于節(jié)點(diǎn)之間所需通信量的增加以及缺乏并行性，「naive implementations」導(dǎo)致運(yùn)行速度急劇下降。
通過(guò)與 FairScale 的緊密合作，現(xiàn)在可以在所有 lightning 模塊上實(shí)現(xiàn) 55% 以上的內(nèi)存減少，只需通過(guò)一個(gè)單一的 flag，這意味著更大的機(jī)型可以適應(yīng)內(nèi)存有限的多張 GPU。
在不更改代碼的情況下啟用 Sharded Training
為了展示在 Lightning 中使用 Sharded Training 有多簡(jiǎn)單，使用 NVIDIA 的一個(gè)流行庫(kù) NeMo 來(lái)訓(xùn)練 Lightning 支持的對(duì)話 AI 模型。使用 NeMo 中提供的 vanilla Transformer LM 模型，有 12 億個(gè)參數(shù)，該模型對(duì)訓(xùn)練內(nèi)存要求很高。在訓(xùn)練大型語(yǔ)言模型時(shí)，內(nèi)存是提高模型大小或提升 GPU 飽和度的寶貴資源。此外使用 WikiText 數(shù)據(jù)集訓(xùn)練模型。
首先下載數(shù)據(jù)集并使用 NVIDIA NeMo 提供的處理腳本進(jìn)行提取，然后在 NeMo 中找到預(yù)配置文件定義模型配置，修改數(shù)據(jù)輸入指向自定義數(shù)據(jù)集。為了進(jìn)行基準(zhǔn)測(cè)試，還構(gòu)建了一個(gè)簡(jiǎn)單的基于單詞的詞匯表。
在設(shè)置模型參數(shù)之后，用戶只需要將 Sharded 插件 flag 傳遞給支持 Sharded Traing 的訓(xùn)練器就可以了。用戶還可以通過(guò)增加 GPU 數(shù)量和啟用本地混合精度（native mixed precision）來(lái)實(shí)現(xiàn)內(nèi)存和速度的進(jìn)一步提升。分區(qū)優(yōu)化器和 GPU 之間的通信可以在后臺(tái)自動(dòng)處理。
下面介紹了使用 Lightning 內(nèi)置 Sharding 與普通 GPU 擴(kuò)展時(shí)每臺(tái)設(shè)備的內(nèi)存提升情況，每臺(tái)設(shè)備的內(nèi)存分配保持不變。不僅如此，Lightning 團(tuán)隊(duì)還給出了 SwAW、DeepSpeech 2 和 iGPT 等其他 PyTorch Lightning 支持模型的測(cè)試結(jié)果。
結(jié)果表明，每個(gè) GPU 上最高節(jié)省內(nèi)存 15GiB，從而可以增加模型能力。例如，在硬件一樣的情況下，訓(xùn)練 Transformer LM 時(shí)模型參數(shù)量可以從 12 億增至 20 億。
使用 8 個(gè) A100s 時(shí)訓(xùn)練 Transformer LM、SwAV Wide ResNet、DeepSpeech2 和 iGPT 時(shí)的平均峰值內(nèi)存比較。
隨著 GPU 之間通信的優(yōu)化，與標(biāo)準(zhǔn)分布式加速器相比，節(jié)點(diǎn)內(nèi)性能的擴(kuò)展效果更好。請(qǐng)注意，隨著向很多節(jié)點(diǎn)的擴(kuò)展，內(nèi)存降低的效果開(kāi)始減弱，這是因?yàn)槠渌蛩爻蔀榱似款i。但是，Sharded training 依然帶來(lái)良好的 throughout 擴(kuò)展。
在 8 個(gè)具有相同超參數(shù)和批大小的 A100s 上的平均 Epoch time 比較，越低越好。
博客地址：https://seannaren.medium.com/introducing-pytorch-lightning-sharded-train-sota-models-with-half-the-memory-7bcc8b4484f
版權(quán)聲明：本文內(nèi)容由互聯(lián)網(wǎng)用戶自發(fā)貢獻(xiàn)，該文觀點(diǎn)僅代表作者本人。本站僅提供信息存儲(chǔ)空間服務(wù)，不擁有所有權(quán)，不承擔(dān)相關(guān)法律責(zé)任。如發(fā)現(xiàn)本站有涉嫌抄襲侵權(quán)/違法違規(guī)的內(nèi)容， 請(qǐng)發(fā)送郵件至 舉報(bào)，一經(jīng)查實(shí)，本站將立刻刪除。
