本文轉自雷鋒網，如需轉載請至雷鋒網官網申請授權。

大型神經網絡是當前人工智能領域的熱門話題之一，那么，如何訓練大模型？

最近，曾推出大規模預訓練模型 GPT-3 的 OpenAI 發表了一篇博文，介紹了基于 GPU 的四種節省內存的并行訓練方法，分別是：

• 數據并行——在不同的 GPU 上運行同一批次的不同子集；
• 流水線并行——在不同的 GPU 上運行模型的不同層；
• 張量并行——分解單個運算的數學運算，例如將矩陣乘法拆分到 GPU 上；
• 專家混合（MOE）——僅通過每層的一小部分處理每個示例。

圖注：三層模型上各種并行策略，每種顏色代表一層，虛線分隔不同的 GPU。

1 數據并行

「數據并行訓練」意味著將相同的參數復制到多個 GPU（通常稱為“workers”），并為每個 GPU 分配不同的示例以同時處理。

單單的數據并行要求模型匹配單個 GPU 內存，但當你利用多個 GPU 計算時，代價是存儲參數的多個副本。不過，話雖如此，有一些策略可以增加 GPU 可用的有效 RAM，例如，在兩次使用之間，可將參數暫時卸載到 CPU 內存。

隨著每次數據并行 worker 更新其參數副本，它們需要相互協調，以確保每個 worker 都繼續具有相似的參數。最簡單的方法是在 worker 之間引入「阻塞通信」：

步驟 1：獨立計算每個worker上的梯度；

步驟 2：將不同 worker 的梯度平均；

步驟 3：在每個 worker 上獨立計算相同的新參數。

步驟 2 是一個阻塞平均值，它需要傳輸大量數據（與 worker 數量乘以參數大小成正比），這可能會損害訓練的吞吐量。有各種異步同步方案可以消除這種損耗，但會損害學習效率；因此在實踐中，人們普遍堅持同步方法。

2 流水線并行

在流水線并行訓練中，研究者會將模型的順序塊劃分到 GPU 上，每個 GPU 只保存一小部分參數，因此，相同模型的每個 GPU 消耗的內存按比例減少。

將大型模型拆分為連續層的塊很簡單，但由于層的輸入和輸出之間存在順序依賴關系，因此，在 worker 等待前一臺機器的輸出用作其輸入時，一個幼稚的執行可能會導致出現大量空閑時間。這些等待時間塊被稱為「泡沫」（bubbles），即浪費了本可以由空閑機器來完成的計算。

圖注：一個簡單的流水線并行設置插圖，其中，模型被垂直分成 4 個分區。worker 1 主持第一層的模型參數（最接近輸入），而 worker 4 主持第 4 層（最接近輸出）。“F”、“B”和“U”分別代表前向、后向和更新操作。下標會指示在哪個 worker 上運行操作。由于順序依賴性，數據一次由一個 worker 處理，導致產生了大量的空閑時間“泡沫”。

我們可以重用數據并行的想法，通過讓每個 worker 一次只處理數據元素的一個子集，來降低產生時間泡沫的成本，從而使我們能巧妙地將新計算與等待時間重疊。核心思想是，將一個批次拆分為多個微批次，每個微批次的處理速度都應該成比例地加快，并且每個 worker 在下一個微批次可用時立即開始工作，從而加快管道執行。有了足夠的微批次， worker 可以在大部分時間被利用，并且在步驟開始和結束時「泡沫」最小。梯度在微批次之間進行平均，并且只有在所有微批次完成后才會更新參數。

模型拆分的 worker 數量通常稱為「管道深度」（pipeline depth）。

在前向傳遞期間，worker 只需將其層塊的輸出（稱為「激活」）發送給下一個 worker；在反向傳遞期間，它僅將這些激活的梯度發送給前一個工作人員。如何安排這些通道以及如何跨微批次聚合梯度有很大的設計空間。例如，方法 GPipe 是讓每個工作進程連續向前和向后傳遞，然后在最后同步聚合來自多個微批次的梯度；而 PipeDream 會安排每個 worker 交替處理的前向和后向通道。

圖注：GPipe 和 PipeDream 流水線方案的比較，每批使用 4 個微批次。微批次 1-8 對應于兩個連續的數據批次。圖中“number”表示在哪個微批次上操作，下標標記 worker  ID。注意，PipeDream 通過使用陳舊參數執行一些計算來獲得更高的效率。

3 張量并行

管道并行性將模型逐層“垂直”拆分，也可以在一個層內“水平”拆分某些操作，這通常稱為張量訓練。

對于許多現代模型（例如Transformer），計算瓶頸是將激活批處理矩陣與大權重矩陣相乘。矩陣乘法可以認為是成對的行和列之間的點積；可以在不同的 GPU 上計算獨立的點積，或者在不同的 GPU 上計算每個點積的部分并總結結果。無論采用哪種策略，我們都可以將權重矩陣分割成大小均勻的“碎片”，將每個碎片托管在不同的 GPU 上，并使用該碎片計算整個矩陣乘積的相關部分，然后再進行通信以組合結果。

一個例子是Megatron-LM，它在 Transformer 的自注意力和 MLP 層內并行化矩陣乘法。PTD-P使用張量、數據和流水線并行，其流水線調度為每個設備分配了多個不連續的層，以增加網絡通信為代價來減少泡沫損耗。

有時，網絡輸入可以跨維度并行化，相對于交叉通信具有高度的并行計算。序列并行就是這樣一種想法，其中輸入序列在時間上被分成多個子示例，通過允許計算繼續進行更細粒度的示例，來按比例減少峰值內存消耗。

4 專家混合 (MoE)

使用專家混合（MoE）方法，只有小部分網絡用于計算任何一個輸入的輸出。

一個示例方法是擁有多組權重，并且網絡可在推理時通過門控機制選擇要使用的權重組，這能在不增加計算成本的情況下啟用更多參數。每組權重都被稱為“專家”，且希望網絡能學會為每個專家分配專門的計算和技能。不同的專家可以主持不同的 GPU ，從而提供了一種明確的方式來擴大用于模型的 GPU 數量。

圖注：門控網絡只選擇了n個專家中的2個。

GShard 將 MoE Transformer 的參數擴展到 6000 億個參數，其中僅將 MoE 層拆分到多個 TPU 設備上，其他層則完全復制。Switch Transformer 通過將一個輸入路由到單個專家，將模型大小擴展到數萬億個參數，具有更高的稀疏性。

5 其他節省內存的設計

還有許多其他的計算策略，可以使訓練越來越大的神經網絡更容易處理。例如：

要計算梯度，需要保存原始激活，這會消耗大量設備 RAM。檢查點（也稱為激活重新計算）存儲激活的任何子集，并在反向傳遞期間，及時重新計算中間的激活，以最多一個額外完整前向傳遞的計算成本，節省了大量內存。人們還可以通過選擇性激活重新計算，來不斷權衡計算和內存成本，這是對激活的子集進行檢查，其存儲成本相對較高，但計算成本較低。

混合精度訓練是使用較低精度的數字（最常見的是FP16）來訓練模型。現代加速器可以使用較低精度的數字達到更高的 FLOP 計數，并且還能節省設備 RAM。在適當的照顧下，產生的模型幾乎可以不損失任何精度。

卸載是將未使用的數據臨時卸載到 CPU 或不同設備之間，在需要時將其讀回。幼稚的執行會大大減慢訓練速度，但復雜的實現方式會預先獲取數據，使設備永遠不需要等待。這個想法的一個實現是ZeRO，它可將參數、梯度和優化器狀態分割到所有可用的硬件上，并根據需要將它們具體化。

Memory Efficient Optimizers已經提出了內存效率優化器，以減少優化器所維護的運行狀態的內存占用，例如Adafactor。

壓縮也可用于存儲網絡中的中間結果。例如，Gist壓縮為后向傳遞而保存的激活；DALL-E在同步梯度之前壓縮梯度。