OpenAI：訓練大型神經網路的四種基本方法

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

大型神經網路是目前人工智慧領域的熱門話題之一，那麼，如何訓練大模型？

最近，曾推出大規模預訓練模型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在同步梯度之前壓縮梯度。

以上是OpenAI：訓練大型神經網路的四種基本方法的詳細內容。更多資訊請關注PHP中文網其他相關文章！