今年 2 月，我們發佈了 DeepSpeed和 ZeRO（Zero Redundancy Optimizer，零冗餘優化器），前者是一個開源的深度學習訓練優化庫，後者是該庫中的一種新型內存優化技術，通過提高規模、速度、成本和可用性，極大地推進了大型模型訓練。DeepSpeed 研究人員創建了Turing-NLG，這是最大的語言模型，擁有 170 億個參數，發佈時具有最先進的正確率。
今年 5 月，我們發佈了 ZeRO-2，它支持 2000 億個參數的模型訓練，比現有技術快 10 倍，同時還提供了一系列計算、I/O 和收斂優化，以實現最快的 BERT 訓練。從那時起，我們一直在快速持續創新，推動深度學習訓練速度和規模的邊界。
今天，我們很高興與大家分享我們的新進展，這些新進展不僅將深度學習訓練推向極致，而且還爲更多的人實現了深度學習訓練的民主化：從在大規模超級計算機上進行訓練的數據科學家，到在低端集羣甚至在單 GPU 上進行訓練的數據科學家。更具體地說，DeepSpeed 增加了四項新的技術，進一步推進了 “AI at Scale” 計劃，以在 Microsoft 的人工智能產品和平臺上進行創新。它們提供了極高的計算、內存和通信效率，併爲模型訓練提供了數十億到數萬億個參數。這些技術還允許在硬件系統上使用極長的輸入序列，這些硬件系統包括單 GPU 的硬件系統、具有數千個 GPU 的高端集羣或具有非常慢的以太網的低端集羣。
- 具有 3D 並行性的萬億參數模型訓練：DeepSpeed 支持三種並行方法的靈活組合——ZeRO 支持的數據並行、流水線並行和張量切片模型並行。3D 並行可適應工作負載需求的變化，爲具有超過一萬億個參數的超大型模型提供支持，同時實現近乎完美的內存擴展和吞吐量擴展效率。此外，其改進的通信效率允許用戶在網絡帶寬有限的常規集羣上將數十億參數模型的訓練速度提高 2~7 倍。
- 使用 ZeRO-Offload，在單 GPU 上進行 10 倍的模型訓練。我們擴展了 ZeRO-2，利用 CPU 和 GPU 內存來訓練大型模型。使用一臺NVIDIA V100 GPU機器，我們的用戶可以在不耗盡內存的情況下運行多達 130 億個參數模型，比現有方法大 10 倍，同時獲得具有競爭力的吞吐量。這一特性實現了數十億參數模型訓練的民主化，爲衆多深度學習從業者打開了探索更大更好模型的窗口。
- 通過 DeepSpeed 稀疏注意力機制，提供 10 倍長序列和 6 倍快的執行速度。DeepSpeed 提供了稀疏注意力內核，這是一種支持長序列模型輸入的工具性技術，無論是文本、圖像還是聲音。與經典的密集 Transformer 相比，它提供了一個數量級更長的輸入序列，並以相當的正確率獲得高達 6 倍的執行速度。它的執行速度也比最先進的稀疏實現快 1.5~3 倍。此外，我們的稀疏內核支持高效執行靈活的稀疏格式，並賦予用戶對其定製的稀疏結構進行創新的能力。
- 1 位 Adam，通信量最多可減少 5 倍。Adam 是訓練許多大規模深度學習模型的有效（且可能是利用最充分的）優化器。然而，Adam 通常與通信高效的優化算法不兼容。因此，在跨分佈式設備擴展時，通信成本可能會成爲瓶頸。我們引入了一種新的算法，具有高效實現的 1 位 Adam，在達到與 Adam 相似的收斂效率的同時，通信量降低了 5 倍。在通信受限的場景中，我們觀察到，分佈式訓練的速度最高可提高 3.5 倍，從而可以擴展到不同類型的 GPU 集羣和網絡。
這篇博文更深入地探討了這四種技術。我們已經在開源庫 DeepSpeed中提供了所有這些令人興奮的新優化。
3D 並行性：擴展到萬億參數模型
隨着現代 GPU 集羣上可用計算的快速增長，訓練一個具有不可思議能力的、強大的萬億參數模型不再是遙不可及的夢想，而是一個近在眼前的現實。DeepSpeed 結合了三種強大的技術，可以訓練萬億規模的模型，並擴展到數千個 GPU：數據並行訓練、模型並行訓練和流水線並行訓練。這種共生關係使深度學習訓練的規模遠遠超出了每種策略單獨提供的能力。3D 並行性同時解決了訓練萬億參數模型的兩個基本挑戰：內存效率和計算效率。因此，DeepSpeed 可以在不犧牲速度的情況下，擴展以適應內存中最大規模的模型。
通過 3D 並行性實現內存和計算效率
數據、模型和流水線並行性各自在提高內存和計算效率方面起着特定的作用。圖 1 說明了我們的 3D 策略。
內存效率：將模型的各層劃分爲流水線階段，並通過模型的並行性進一步劃分每個階段的層。這種 2D 組合同時降低了模型、優化器和激活所消耗的內存。然而，我們並不能無限期地對模型進行劃分而不屈服通信開銷，從而限制了計算效率。
計算效率：爲了在不犧牲計算效率的情況下，允許工作器（worker）的數量超過模型和流水線的並行性，我們使用了 ZeRO 支持的數據並行性（ZeRO-DP）。ZeRO-DP 不僅可以通過優化器狀態分區進一步提高內存效率，而且還可以通過拓撲感知映射以最小的通信開銷擴展到任意數量的 GPU。
拓撲感知的 3D 映射（圖 2）：3D 並行中的每個維都被仔細地映射到工作器上，通過利用兩個關鍵的架構屬性來實現最大的計算效率。
- 優化節點內和節點間的通信帶寬：在三種策略中，模型並行的通信開銷最大，因此我們優先將模型並行組放置在節點內以利用更大的節點內帶寬。在這裏，我們應用 NVIDIA Megatron-LM 進行張量分割風格的模型並行。當模型並行不跨越節點中的所有工作器時，則將數據並行組放置在節點內；否則，將數據並行組放置在節點間；流水線並行具有最低的通信量，因此我們可以在不受通信帶寬限制的情況下跨節點調度流水線階段。
- 在通信中通過並行實現帶寬放大：通過流水線並行和模型並行，每個數據並行組通信的梯度大小線性下降，從而使總通信量從純數據並行減少，並且每個數據並行組在局部工作器子集之間獨立並行地執行通信。結果，通過減少通信量、增加局部性和並行性的結合，數據並行通信的有效帶寬被放大。
每個流水線階段內的層在四個模型並行工作器之間進一步劃分。最後，每個流水線跨兩個數據並行實例進行復制，ZeRO 將優化器狀態跨數據並行副本進行分區。
利用線性效率擴展支持萬億級參數模型訓練
DeepSpeed 可以使用至少800 個 NVIDIA V100 GPU 來訓練一個具有一萬億參數的語言模型（圖 3）。我們通過擴展模型的規模，並觀察到線性增長，無論是在模型的規模還是訓練的吞吐量方面，都展示了同步的內存和計算效率。在每個配置中，我們每個 GPU 可以訓練大約 14 億個參數，這是單個 GPU 可以支持的最大模型大小，而不會耗盡內存，這表明了完美的內存擴展。我們還獲得了接近完美的線性計算效率擴展和每個 V100 GPU 47 每秒萬億次浮點運算的吞吐量。對於給定的硬件來說，這是令人印象深刻的可擴展性和吞吐量。
DeepSpeed 可以使用 800 個具有 32 GB 內存的 NVIDIA V100 張量核心 GPU 來訓練具有 1 萬億個參數的模型。每種配置都使用由NVIDIA Megatron-LM提供的 16 路模型並行，其餘 GPU 使用流水線並行進行排列。萬億參數模型有 298 層 Transformer，隱維度爲 17408，用序列長度爲 2048，批大小爲 2048。對於較小的模型，我們按 GPU 數量的比例減少 Transformer 層數和批處理大小。
ZeRO-Offload：使用單 GPU 將模型訓練擴大 10 倍
ZeRO-Offload 通過利用 GPU 及其主機 CPU 上的計算和內存資源，推入最大模型大小的邊界，從而可以使用最小的 GPU 資源有效地訓練。它允許在單個 NVIDIA V100 GPU 上訓練多達 130 億個參數模型，比最先進的 GPU 大 10 倍，同時保持每個 GPU 超過 30 每秒萬億次浮點運算的高訓練吞吐量。
通過在單個 GPU 上實現數十億參數模型訓練，ZeRO-Offload 使大型模型訓練民主化，使資源有限的深度學習從業者也可以使用大型模型。
ZeRO-Offload 背後的關鍵技術是我們在ZeRO-2的基礎上，將優化器狀態和梯度卸載到 CPU 內存上的新功能。這種方法允許 ZeRO-Offload 最小化 CPU 卸載帶來的計算效率損失，同時還可以實現與原始 ZeRO-2 相同的、有時甚至更好的效率。下圖顯示了 ZeRO-Offload 的架構。
DeepSpeed 稀疏注意力機制：爲 10 倍長的序列提供支持，並將執行速度提高 6 倍
基於注意力的深度學習模型，比如 Transformer，在捕獲輸入序列中令牌之間的關係方面非常有效，甚至可以跨越很長的距離。因此，它們與文本、圖像和基於聲音的輸入一起使用，其中序列長度可以用數千個令牌表示。然而，儘管注意力模塊能夠有效地捕捉長期的依賴關係，但在實踐中，它們對長序列輸入的應用受到注意力計算的計算量和記憶需求的限制，這些計算量和內存需求隨着序列長度 $n$ 的增長而呈二次方增長 $O(n^2)$。
爲了解決這一侷限性，DeepSpeed 提供了一套稀疏注意力機制，這是一種工具性技術，可以通過塊稀疏計算以數量級減少注意力計算的計算和內存需求。該套件不僅緩解了注意力計算的內存瓶頸，而且能有效地進行稀疏計算。它的 API 允許與任何基於 Transformer 的模型進行方便的集成。除了提供廣泛的稀疏結構之外，它還具有處理任何用戶定義的塊稀疏結構的靈活性。
更具體地說，稀疏注意力可以被設計成計算附近令牌之間的局部注意力，或者經由利用局部注意力計算的彙總令牌來計算全局注意力。此外，稀疏注意力還可以允許隨機注意力或局部注意力，全局注意力和隨機注意力的任意組合，如圖 6 所示，分別使用藍色、橙色和綠色塊。結果，稀疏注意力將內存佔用減少爲 $O(wn)$，其中 $1<w\leq n$ 是一個參數，其值取決於注意力結構。
GPU 的有效實現：雖然稀疏注意力的基本實現可能顯示出節省內存的好處，但在計算方面來看，它可能甚至比完全計算更糟糕。這主要是由於稀疏數據給全圖增加了分歧和未合併的內存訪問。一般來說，開發高效的稀疏內核，特別是在 GPU 上，是一項具有挑戰性的工作。DeepSpeed 提供了Triton開發的高效稀疏注意力內核。這些內核是以塊稀疏範式結構的，它支持對齊內存訪問，緩解線程分歧，並平衡處理器上的工作負載。
系統性能：稀疏注意力支持10 倍以上的序列，計算速度快達 6.3 倍以上，如圖 7 所示。左圖顯示了三種設置下 BERT-Base 和 BERT-Large 模型中可運行的最長序列長度：密集、帶激活檢查點的密集和帶激活檢查點的稀疏。對於 BERT-Base 和 BERT-Large，與密集註意力相比，稀疏注意力算法減少了總計算量，提高了訓練速度：隨着序列長度的增加，其 boost 值增大，而 BERT-Base 和 BERT-Large 分別提高了 6.3 倍和 5.3 倍。
靈活處理任何塊狀稀疏結構：DeepSpeed 稀疏注意力套件並不針對任何特定的稀疏結構，而是讓模型科學家能夠在高效的系統支持下探索任何塊狀稀疏結構。目前，我們已經添加了流行的稀疏結構，如Fixed（來自 OpenAI Sparse Transformer）、BigBird（來自 Google）和 BSLongformer（AI2Longformer的 Block-Sparse 實現）。我們還定義了一個具有 "變量 "結構的模板，如圖 6 所示，它可以用來簡單地定製任何塊-稀疏的隨機、局部或全局注意力模式。
1 位 Adam：通信量減少 5 倍，訓練速度提高 3.4 倍
大型模型（如 BERT 和 GPT-3）的可擴展訓練需要根據模型設計、架構和系統功能進行仔細的優化。從系統的角度來看，通信已成爲主要瓶頸，特別是在具有標準 TCP 互連的商用系統上，這些互連提供了有限的網絡帶寬。
通信壓縮是減少此類系統訓練時間的一項重要技術。壓縮通信最有效的方法之一是通過誤差補償壓縮，它提供了強大的收斂速度，即使在 1 位壓縮下也是如此。然而，最先進的錯誤補償技術只適用於基本的優化器，如隨機梯度下降（SGD）和動量 SGD，它們是線性依賴梯度的。它們不能與 Adam 等基於梯度的非線性優化器一起工作，Adam 爲許多任務提供了最先進的收斂效率和正確率，包括訓練類 BERT 的模型。
對於像 Adam 這樣強大的優化器來說，由於梯度（在方差項中）的非線性依賴性，使得開發基於誤差補償的壓縮技術具有挑戰性，限制了最先進的通信壓縮技術的實用價值。
用 1 位 Adam 壓縮通信
爲了在使用 Adam 優化器的同時壓縮通信，我們開發了1 位 Adam，它通過預處理來解決梯度中的非線性問題。我們觀察到非線性項方差（$v_t$）的變化幅度在經過幾個輪數的訓練後顯著減小，並且在之後設置 $v_t$ 常數不會改變收斂速度。建議的 1 位 Adam 優化器，如圖 14 所示，由兩部分組成：預熱階段，本質上是普通 Adam 算法；壓縮階段，保持方差項不變，並將剩餘的線性項，即動量，壓縮爲 1 位表示。
該算法的壓縮階段由一個閾值參數控制（如圖 8 所示）。當我們檢測到“方差”的變化低於某個閾值時，我們就切換到壓縮階段。我們的研究表明，預熱階段只需整個訓練步驟的 15%~20%。
解決 1 位 Adam 的系統挑戰
除了算法挑戰之外，在訓練系統中應用 1 位 Adam 還有兩個系統挑戰。首先，我們需要將動量轉換爲 1 位表示的高效內核。第二，我們需要高效的通信方案來跨不同的 GPU 交換這種壓縮的動量。壓縮的目標是減少整體訓練時間，以便具有帶寬受限互聯的商品系統可以用來訓練大型模型。我們在 DeepSpeed 中解決了這些挑戰，並引入了一個完全優化的 1 位 Adam 實現，用於在通信受限的系統上進行訓練。
1 位 Adam 在通信受限系統中的優勢
1 位 Adam 提供了與 Adam 相同的收斂性，產生的通信量減少了 5 倍，使得 BERT-Large 預訓練的吞吐量提高了 3.5 倍，SQuAD 微調的吞吐量提高了 2.7 倍。在壓縮階段觀察到 6.6 倍（圖 9 左）和 6.2 倍（圖 9 右）的速度提升，使得這種端到端的吞吐量提升得以實現。值得一提的是，我們的 1 位 Adam 優化器在 40 千兆以太網系統上的擴展性非常好，其性能與 Adam 在 40 千兆 InfiniBand QDR 系統上的擴展性相當。我們注意到，根據 iPerf 基準，40 千兆以太網上的有效帶寬爲 4.1 Gbps，而根據 InfiniBand perftest 微基準，InfiniBand 提供了接近峯值的 32 Gbps 帶寬。
有關這些新技術的代碼、教程和文檔，請訪問DeepSpeed 網站和GitHub 倉庫。我們還將其中一些技術集成到ONNX 運行時中。
作者介紹：
Rangan Majumder，Microsoft 負責搜索和人工智能的副總裁。
Junhua Wang，WebXT 搜索與人工智能平臺團隊的副總裁、傑出工程師。
原文鏈接：
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/