本文轉(zhuǎn)自雷鋒網(wǎng)，如需轉(zhuǎn)載請至雷鋒網(wǎng)官網(wǎng)申請授權(quán)。

大型神經(jīng)網(wǎng)絡(luò)是當(dāng)前人工智能領(lǐng)域的熱門話題之一，那么，如何訓(xùn)練大模型？

最近，曾推出大規(guī)模預(yù)訓(xùn)練模型 GPT-3 的 OpenAI 發(fā)表了一篇博文，介紹了基于 GPU 的四種節(jié)省內(nèi)存的并行訓(xùn)練方法，分別是：

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

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

1 數(shù)據(jù)并行

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

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

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

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

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

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

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

2 流水線并行

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

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

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

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

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

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

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

3 張量并行

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

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

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

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

4 專家混合 (MoE)

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

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

圖注：門控網(wǎng)絡(luò)只選擇了n個專家中的2個。

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

5 其他節(jié)省內(nèi)存的設(shè)計

還有許多其他的計算策略，可以使訓(xùn)練越來越大的神經(jīng)網(wǎng)絡(luò)更容易處理。例如：

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

混合精度訓(xùn)練是使用較低精度的數(shù)字（最常見的是FP16）來訓(xùn)練模型?，F(xiàn)代加速器可以使用較低精度的數(shù)字達(dá)到更高的 FLOP 計數(shù)，并且還能節(jié)省設(shè)備 RAM。在適當(dāng)?shù)恼疹櫹拢a(chǎn)生的模型幾乎可以不損失任何精度。

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

Memory Efficient Optimizers已經(jīng)提出了內(nèi)存效率優(yōu)化器，以減少優(yōu)化器所維護(hù)的運行狀態(tài)的內(nèi)存占用，例如Adafactor。

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