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