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

語言: CN / TW / HK

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

參考連結:

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

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

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

「其他文章」