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