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