作者 | 包永剛、陳大鑫
微軟於近日釋出了DeepSpeed庫的更新版本，該庫引入了一種新方法來訓練包含 數萬億個引數的AI模型，這種模型內部的變數可為預測提供依據。微軟稱這種技術稱為3D並行性，可以適應各種的工作負載需求，特別是為功耗大的超大型模型平衡效率。
具有數十億個引數的單個大型AI模型在一系列具有挑戰性的領域中取得了長足的進步。研究表明它們之所以表現出色，是因為它們可以理解語言、語法、知識、概念和上下文的細微差別，使他們能夠總結演講，在實時遊戲聊天中過濾不合適的話語，解析複雜的法律文件，甚至可以透過搜尋GitHub生成程式碼。
但是訓練模型需要大量的計算資源。根據2018年OpenAI的分析，從2012年到2018年，在大型AI訓練中，計算量的需求增長了300,000倍，也就是大概每3.5個月的時間翻一番，遠遠超過了摩爾定律的步伐。
增強的DeepSpeed利用三種技術來實現“萬億級”模型訓練：資料並行訓練、模型並行訓練和pipeline並行訓練。
訓練一個萬億引數的模型將需要至少400個Nvidia 今年最新推出的A100 GPU（每個記憶體高達40GB），並且Microsoft估計需要4,000臺A100以50%的效率執行，大約需要100天才能完成訓練。這與微軟和OpenAI共同設計的AI超級計算機無並不匹配，後者包含10,000多個顯示卡，在這樣的超大規模下，很難獲得很高的計算效率。
3D並行性中的每個維度都被仔細地對映到workers上，以透過利用兩個關鍵的架構屬性來實現最大的計算效率。
1、最佳化節點內和節點間的通訊頻寬：模型並行性是這三種策略中最大的通訊開銷，因此，需要優先考慮將模型並行組放置在節點內以利用更大的節點內頻寬。在這裡，將NVIDIA Megatron-LM應用為模型並行度的張量切片樣式。當模型並行性不覆蓋節點中的所有工作執行緒時，資料並行組將放置在節點內。否則，它們將跨節點放置。流水線並行性具有最低的通訊量，因此可以跨節點排程流水線階段，而不受通訊頻寬的限制。
2、透過通訊中的並行性進行頻寬放大：每個資料並行組透過流水線和模型並行性通訊的梯度大小線性減小，因此總通訊量從純資料並行性開始減少。此外，每個資料並行組在區域性工作程式的子集之間獨立且並行地執行其通訊。最後透過減少通訊量以及增加區域性性和並行性的組合來放大用於資料並行。
DeepSpeed將大型模型分為四個流水線階段，分為較小的元件（層）。
每個流水線階段中的層進一步劃分為四個“workers”，它們執行實際的訓練。每個管道都在兩個並行資料例項之間複製，並且工作程式被對映到具有八個節點（每個節點具有四個GPU）的系統上的GPU，彩色表示同一節點上的GPU。。
由於有了這些和其他效能改進，微軟表示，一個萬億引數的AI模型可以減少到僅需要800個Nvidia V100 GPU訓練。
DeepSpeed的最新版本還包含了ZeRO-Offload技術，透過利用GPU及其主機CPU上的計算和記憶體資源，ZeRO-Offload突破了可以使用最少GPU資源進行有效訓練的最大模型尺寸的邊界。它可以在單個NVIDIA V100 GPU上訓練多達130億個引數的模型，比最新技術大10倍，同時保持每個GPU超過30 teraflops的高訓練吞吐量。
圖注：可以在單個GPU上使用預設的PyTorch和ZeRO-Offload訓練最大的模型。
透過在單個GPU上啟用數十億引數的模型訓練，ZeRO-Offload使大型模型訓練平民化，從而使資源有限的深度學習從業人員可以使用它。
“這些（DeepSpeed中的新技術）提供了極高的計算、記憶體和通訊效率，並支援數十億到數萬億個引數的模型訓練。” 微軟在部落格中寫道，“這些技術還允許極長的輸入序列，並透過單個GPU，具有數千個GPU的高階叢集或具有非常慢的乙太網網路的低端叢集釋放硬體系統的能量。我們將（繼續）進行快速創新，為深度學習訓練突破了速度和規模的界限。”
編譯via venturebeat、
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/