來源：新智元
本文約3000字，建議閱讀5分鐘
本文詳細介紹了一些訓練大型神經(jīng)網(wǎng)絡的相關技術及底層原理。

想知道那些超大規(guī)模神經(jīng)網(wǎng)絡都是怎么訓出來的？OpenAI一篇文章總結：除了顯卡要多，算法也很重要！

如今AI的很多進步都要歸功于大型神經(jīng)網(wǎng)絡，尤其是大公司和研究機構提供的預訓練模型更是推動了下游任務的進步。

但想自己動手訓練一個大型神經(jīng)網(wǎng)絡并不簡單，首先要面對的就是海量的數(shù)據(jù)、多機協(xié)調和大量GPU的調度工作。

一提到「并行」，冥冥之中就會感覺多了很多隱藏的bug。

最近OpenAI發(fā)布了一篇文章，詳細介紹了一些訓練大型神經(jīng)網(wǎng)絡的相關技術及底層原理，徹底消除你對并行的恐懼！

比如以并行訓練一個三層的神經(jīng)網(wǎng)絡為例，其中并行可以分為數(shù)據(jù)并行、pipeline并行，trensor并行和專家并行，圖中不同顏色代表不同層、虛線隔開的是不同的GPU。

聽上去很多，但理解這些并行技術實際上只需要對計算結構進行一些假設，然后對數(shù)據(jù)包的流動方向有所了解即可。

訓練流程無并行

訓練一個神經(jīng)網(wǎng)絡是一個迭代的過程。

在一次迭代中，輸入數(shù)據(jù)經(jīng)過模型的層，前向傳遞后即可為一個batch數(shù)據(jù)中的每個訓練實例計算輸出。

然后各層再向后傳遞，通過計算每個參數(shù)的梯度來傳播每個參數(shù)對最終輸出的影響程度。

每個batch數(shù)據(jù)的平均梯度、參數(shù)和一些每個參數(shù)的優(yōu)化狀態(tài)被傳遞給一個優(yōu)化算法，比如Adam可以計算下一個迭代的參數(shù)（在你的數(shù)據(jù)上應該有更好的性能）和新的每個參數(shù)的優(yōu)化狀態(tài)。

隨著訓練在大量數(shù)據(jù)上的迭代，模型不斷進化，產(chǎn)生越來越精確的輸出。

在整個訓練過程中，會有不同的并行技術在不同的維度上進行切割，包括：

1、數(shù)據(jù)并行，即在不同的GPU上運行一個batch的不同子集；

2、pipeline并行，即在不同的GPU上運行模型的不同層；

3、tensor并行，即將單一操作（如矩陣乘法）的數(shù)學運算拆分到不同的GPU上；

4、專家混合（Mixture of Experts, MoE），即只用每層的一部分來處理每個輸入實例。

并行中說的GPU并非僅局限于GPU，對于其他神經(jīng)網(wǎng)絡加速器的用戶來說，這些想法同樣有效。

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

數(shù)據(jù)并行訓練意味著將相同的參數(shù)復制到多個GPU（通常稱為worker），并將不同的實例分配給每個GPU同時進行處理。

單純的數(shù)據(jù)并行仍然需要模型符合單個GPU的內存要求，如果你利用多個GPU進行計算，那代價就是存儲許多重復的參數(shù)副本。有一些策略可以增加你的GPU可用的有效RAM，比如在兩次使用之間將參數(shù)暫時卸載到CPU內存。

當每個數(shù)據(jù)并行worker更新其參數(shù)副本時，他們需要協(xié)調以確保每個worker繼續(xù)擁有類似的參數(shù)。

最簡單的方法是在worker之間引入阻塞式通信：

1、在每個worker上獨立計算梯度；

2、在各worker上平均梯度；

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

其中步驟2是的阻塞需要傳輸相當多的數(shù)據(jù)（與worker的數(shù)量乘以參數(shù)量的大小成正比），非常有可能降低訓練吞吐量。

雖然有各種異步同步方案來消除這種開銷，但它們會損害學習效率；在實踐中，研究人員通常還是會堅持使用同步方法。

Pipeline并行

pipeline并行訓練的意思是將模型的順序塊分割到不同的GPU上，每個GPU只持有一部分參數(shù)，因此，同一個模型在每個GPU上消耗的內存比例較小。

將一個大的模型分割成連續(xù)層的大塊是很直接的一種方式。然而，各層的輸入和輸出之間存在著順序上的依賴性，所以一個樸素的實現(xiàn)可能會導致大量的空閑時間，而wroker在等待前一個機器的輸出被用作其輸入。

這些等待時間塊被稱為氣泡（bubbles），浪費了空閑機器可以完成的計算。

我們可以重用數(shù)據(jù)并行的思想，通過讓每個worker一次只處理一個數(shù)據(jù)元素的子集來降低氣泡的成本，巧妙地將新的計算與等待時間重疊起來。

核心思想是將一個batch分成多個microbatches；每個微批的處理速度應該是成比例的，每個worker在下一個微批可用時就開始工作，從而加速管道的執(zhí)行。

有了足夠的微批，worker在大部分時間內都處于工作狀態(tài)，并且在每個step的開始和結束時氣泡最小。梯度是所有微批的平均值，只有所有微批完成之后才會進行參數(shù)更新。

模型被分割的worker的數(shù)量通常被稱為pipeline深度。

在前向傳遞期間，worker只需要將其大塊層的輸出（也叫激活）發(fā)送給下一個worker；在后向傳遞期間，它只將這些激活的梯度發(fā)送給前一個worker。如何調度這些傳遞過程以及如何在微批中聚合梯度，仍然有很大的設計空間。

GPipe的做法是讓每個worker連續(xù)地處理前向和后向的傳遞，然后在最后同步地聚合來自多個微批的梯度。而PipeDream則安排每個工作者交替地處理前向和后向通道。

Tensor并行

Pipeline并行是將一個模型按層「垂直」分割，而Tensor并行則是在一個層內「橫向」分割某些操作。

對于現(xiàn)代模型（如Transformer）來說，計算瓶頸主要來自激活批矩陣與大權重矩陣相乘。矩陣乘法可以被認為是成對的行和列之間的點積，所以是有可能在不同的GPU上獨立計算點積，或者在不同的GPU上計算每個點積的一部分，最后再將結果相加。無論采用哪種策略，我們都可以將權重矩陣切成偶數(shù)大小的「碎片」，將每個碎片放在不同的GPU上，并使用該碎片來計算整個矩陣乘積的相關部分，然后再進行GPU間通信來合并結果。

Megatron-LM采用的就是這種方式，它在Transformer的自注意力和MLP層內并行化矩陣乘法。PTD-P使用Tensor、數(shù)據(jù)和Pipeline并行；它的pipeline調度將多個不連續(xù)的層分配給每個設備，以更多的網(wǎng)絡通信為代價減少氣泡的開銷。

有時網(wǎng)絡的輸入也可以在一個維度上進行并行化，相對于交叉通信來說，并行計算的程度很高。序列并行就是這樣一個想法，一個輸入序列在不同時間被分割成多個子實例，通過以更細粒度的實例進行計算，可以按比例減少峰值內存消耗。

混合專家系統(tǒng)（MoE）

對于任何一個輸入，MoE策略都只有一部分網(wǎng)絡被用來計算輸出。

比如說一個網(wǎng)絡里有很多套權重，網(wǎng)絡可以在推理時通過門控機制選擇具體使用哪套。這樣就可以在不增加計算成本的情況下增加參數(shù)量。

每組權重被稱為一個「專家」，訓練目標是希望網(wǎng)絡能夠學會將專門的計算和技能分配給每個專家。不同的專家可以托管在不同的GPU上，為擴大模型使用的GPU數(shù)量提供了一個明確的方法。

GShard可以將MoE Transformer的規(guī)模擴大到6000億個參數(shù)，其方案是只有MoE層被分割到多個TPU設備上，而其他層則是完全重復的。

Switch Transformer則是通過將一個輸入路由（routing）到一個專家，成功將模型規(guī)模擴展到數(shù)萬億的參數(shù)，甚至稀疏度更高。

省內存小妙招

除了買GPU外，還有一些計算策略可以幫助節(jié)省內存，方便訓練更大的神經(jīng)網(wǎng)絡。

1、為了計算梯度，你可能需要保存原始激活值，這可能會消耗大量的設備內存。檢查點（Checkpointing, 也被稱為激活再計算）可以存儲任何激活子集，并在后向通道中以just-in-time的方式重新計算中間激活。

這種方式可以節(jié)省大量的內存，而計算成本最多就是多出一個完整的前向傳遞。我們也可以通過選擇性的激活再計算來不斷地在計算和內存成本之間進行權衡，也就是檢查那些存儲成本相對較高但計算成本較低的激活子集。

2、混合精度訓練（Mixed Precision Training）是使用較低精度的數(shù)字（最常見的是FP16）來訓練模型。現(xiàn)代的計算加速器可以用低精度數(shù)字達到更高的FLOP數(shù)，而且你還可以節(jié)省設備RAM。只要處理得當，這種方式訓練得到的模型在性能上幾乎不會有太大損失。

3、卸載（Offloading）是將未使用的數(shù)據(jù)暫時卸載到CPU或不同的設備中，然后在需要時再將其讀回。樸素的實現(xiàn)方式會大大降低訓練速度，但復雜的實現(xiàn)方式會預先獲取數(shù)據(jù)，這樣設備就不需要再等待了。

這個想法的一個具體實現(xiàn)是ZeRO，它將參數(shù)、梯度和優(yōu)化器狀態(tài)分割到所有可用的硬件上，并根據(jù)實際需要再將它們具體化。

4、內存效率優(yōu)化器（Memory Efficient Optimizer）可以減少優(yōu)化器所維護的運行狀態(tài)的內存占用，如Adafactor。

5、壓縮（Compression）也可用于存儲網(wǎng)絡中的中間結果。例如，Gist對為后向傳遞而保存的激活進行壓縮；DALL-E在同步梯度之前壓縮了梯度。

參考資料：

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

編輯：王菁

校對：龔力

OpenAI煉丹秘籍：教你學會訓練大型神經(jīng)網(wǎng)絡