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

來源 | 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在同步梯度之前壓縮梯度。

參考連結：

http://openai.com/blog/techniques-for-training-large-neural-networks/

雷峰網 (公眾號：雷峰網)

雷峰網版權文章，未經授權禁止轉載。詳情見 轉載須知 。