編輯：LRS
【新智元導(dǎo)讀】想知道那些超大規(guī)模神經(jīng)網(wǎng)絡(luò)都是怎么訓(xùn)出來的？OpenAI一篇文章總結(jié)：除了顯卡要多，算法也很重要！
如今AI的很多進(jìn)步都要歸功于大型神經(jīng)網(wǎng)絡(luò)，尤其是大公司和研究機構(gòu)提供的預(yù)訓(xùn)練模型更是推動了下游任務(wù)的進(jìn)步。但想自己動手訓(xùn)練一個大型神經(jīng)網(wǎng)絡(luò)并不簡單，首先要面對的就是海量的數(shù)據(jù)、多機協(xié)調(diào)和大量GPU的調(diào)度工作。一提到「并行」，冥冥之中就會感覺多了很多隱藏的bug。最近OpenAI發(fā)布了一篇文章，詳細(xì)介紹了一些訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)的相關(guān)技術(shù)及底層原理，徹底消除你對并行的恐懼！比如以并行訓(xùn)練一個三層的神經(jīng)網(wǎng)絡(luò)為例，其中并行可以分為數(shù)據(jù)并行、pipeline并行，trensor并行和專家并行，圖中不同顏色代表不同層、虛線隔開的是不同的GPU。聽上去很多，但理解這些并行技術(shù)實際上只需要對計算結(jié)構(gòu)進(jìn)行一些假設(shè)，然后對數(shù)據(jù)包的流動方向有所了解即可。訓(xùn)練一個神經(jīng)網(wǎng)絡(luò)是一個迭代的過程。在一次迭代中，輸入數(shù)據(jù)經(jīng)過模型的層，前向傳遞后即可為一個batch數(shù)據(jù)中的每個訓(xùn)練實例計算輸出。然后各層再向后傳遞，通過計算每個參數(shù)的梯度來傳播每個參數(shù)對最終輸出的影響程度。每個batch數(shù)據(jù)的平均梯度、參數(shù)和一些每個參數(shù)的優(yōu)化狀態(tài)被傳遞給一個優(yōu)化算法，比如Adam可以計算下一個迭代的參數(shù)（在你的數(shù)據(jù)上應(yīng)該有更好的性能）和新的每個參數(shù)的優(yōu)化狀態(tài)。隨著訓(xùn)練在大量數(shù)據(jù)上的迭代，模型不斷進(jìn)化，產(chǎn)生越來越精確的輸出。在整個訓(xùn)練過程中，會有不同的并行技術(shù)在不同的維度上進(jìn)行切割，包括：1、數(shù)據(jù)并行，即在不同的GPU上運行一個batch的不同子集；2、pipeline并行，即在不同的GPU上運行模型的不同層；3、tensor并行，即將單一操作（如矩陣乘法）的數(shù)學(xué)運算拆分到不同的GPU上；4、專家混合（Mixture of Experts, MoE），即只用每層的一部分來處理每個輸入實例。并行中說的GPU并非僅局限于GPU，對于其他神經(jīng)網(wǎng)絡(luò)加速器的用戶來說，這些想法同樣有效。數(shù)據(jù)并行訓(xùn)練意味著將相同的參數(shù)復(fù)制到多個GPU（通常稱為worker），并將不同的實例分配給每個GPU同時進(jìn)行處理。單純的數(shù)據(jù)并行仍然需要模型符合單個GPU的內(nèi)存要求，如果你利用多個GPU進(jìn)行計算，那代價就是存儲許多重復(fù)的參數(shù)副本。有一些策略可以增加你的GPU可用的有效RAM，比如在兩次使用之間將參數(shù)暫時卸載到CPU內(nèi)存。當(dāng)每個數(shù)據(jù)并行worker更新其參數(shù)副本時，他們需要協(xié)調(diào)以確保每個worker繼續(xù)擁有類似的參數(shù)。3、每個worker上獨立計算相同的新參數(shù)。其中步驟2是的阻塞需要傳輸相當(dāng)多的數(shù)據(jù)（與worker的數(shù)量乘以參數(shù)量的大小成正比），非常有可能降低訓(xùn)練吞吐量。雖然有各種異步同步方案來消除這種開銷，但它們會損害學(xué)習(xí)效率；在實踐中，研究人員通常還是會堅持使用同步方法。pipeline并行訓(xùn)練的意思是將模型的順序塊分割到不同的GPU上，每個GPU只持有一部分參數(shù)，因此，同一個模型在每個GPU上消耗的內(nèi)存比例較小。將一個大的模型分割成連續(xù)層的大塊是很直接的一種方式。然而，各層的輸入和輸出之間存在著順序上的依賴性，所以一個樸素的實現(xiàn)可能會導(dǎo)致大量的空閑時間，而wroker在等待前一個機器的輸出被用作其輸入。這些等待時間塊被稱為氣泡（bubbles），浪費了空閑機器可以完成的計算。我們可以重用數(shù)據(jù)并行的思想，通過讓每個worker一次只處理一個數(shù)據(jù)元素的子集來降低氣泡的成本，巧妙地將新的計算與等待時間重疊起來。核心思想是將一個batch分成多個microbatches；每個微批的處理速度應(yīng)該是成比例的，每個worker在下一個微批可用時就開始工作，從而加速管道的執(zhí)行。有了足夠的微批，worker在大部分時間內(nèi)都處于工作狀態(tài)，并且在每個step的開始和結(jié)束時氣泡最小。梯度是所有微批的平均值，只有所有微批完成之后才會進(jìn)行參數(shù)更新。模型被分割的worker的數(shù)量通常被稱為pipeline深度。在前向傳遞期間，worker只需要將其大塊層的輸出（也叫激活）發(fā)送給下一個worker；在后向傳遞期間，它只將這些激活的梯度發(fā)送給前一個worker。如何調(diào)度這些傳遞過程以及如何在微批中聚合梯度，仍然有很大的設(shè)計空間。GPipe的做法是讓每個worker連續(xù)地處理前向和后向的傳遞，然后在最后同步地聚合來自多個微批的梯度。而PipeDream則安排每個工作者交替地處理前向和后向通道。Pipeline并行是將一個模型按層「垂直」分割，而Tensor并行則是在一個層內(nèi)「橫向」分割某些操作。對于現(xiàn)代模型（如Transformer）來說，計算瓶頸主要來自激活批矩陣與大權(quán)重矩陣相乘。矩陣乘法可以被認(rèn)為是成對的行和列之間的點積，所以是有可能在不同的GPU上獨立計算點積，或者在不同的GPU上計算每個點積的一部分，最后再將結(jié)果相加。無論采用哪種策略，我們都可以將權(quán)重矩陣切成偶數(shù)大小的「碎片」，將每個碎片放在不同的GPU上，并使用該碎片來計算整個矩陣乘積的相關(guān)部分，然后再進(jìn)行GPU間通信來合并結(jié)果。Megatron-LM采用的就是這種方式，它在Transformer的自注意力和MLP層內(nèi)并行化矩陣乘法。PTD-P使用Tensor、數(shù)據(jù)和Pipeline并行；它的pipeline調(diào)度將多個不連續(xù)的層分配給每個設(shè)備，以更多的網(wǎng)絡(luò)通信為代價減少氣泡的開銷。有時網(wǎng)絡(luò)的輸入也可以在一個維度上進(jìn)行并行化，相對于交叉通信來說，并行計算的程度很高。序列并行就是這樣一個想法，一個輸入序列在不同時間被分割成多個子實例，通過以更細(xì)粒度的實例進(jìn)行計算，可以按比例減少峰值內(nèi)存消耗。對于任何一個輸入，MoE策略都只有一部分網(wǎng)絡(luò)被用來計算輸出。比如說一個網(wǎng)絡(luò)里有很多套權(quán)重，網(wǎng)絡(luò)可以在推理時通過門控機制選擇具體使用哪套。這樣就可以在不增加計算成本的情況下增加參數(shù)量。每組權(quán)重被稱為一個「專家」，訓(xùn)練目標(biāo)是希望網(wǎng)絡(luò)能夠?qū)W會將專門的計算和技能分配給每個專家。不同的專家可以托管在不同的GPU上，為擴大模型使用的GPU數(shù)量提供了一個明確的方法。GShard可以將MoE Transformer的規(guī)模擴大到6000億個參數(shù)，其方案是只有MoE層被分割到多個TPU設(shè)備上，而其他層則是完全重復(fù)的。Switch Transformer則是通過將一個輸入路由（routing）到一個專家，成功將模型規(guī)模擴展到數(shù)萬億的參數(shù)，甚至稀疏度更高。除了買GPU外，還有一些計算策略可以幫助節(jié)省內(nèi)存，方便訓(xùn)練更大的神經(jīng)網(wǎng)絡(luò)。1、為了計算梯度，你可能需要保存原始激活值，這可能會消耗大量的設(shè)備內(nèi)存。檢查點（Checkpointing, 也被稱為激活再計算）可以存儲任何激活子集，并在后向通道中以just-in-time的方式重新計算中間激活。這種方式可以節(jié)省大量的內(nèi)存，而計算成本最多就是多出一個完整的前向傳遞。我們也可以通過選擇性的激活再計算來不斷地在計算和內(nèi)存成本之間進(jìn)行權(quán)衡，也就是檢查那些存儲成本相對較高但計算成本較低的激活子集。2、混合精度訓(xùn)練（Mixed Precision Training）是使用較低精度的數(shù)字（最常見的是FP16）來訓(xùn)練模型。現(xiàn)代的計算加速器可以用低精度數(shù)字達(dá)到更高的FLOP數(shù)，而且你還可以節(jié)省設(shè)備RAM。只要處理得當(dāng)，這種方式訓(xùn)練得到的模型在性能上幾乎不會有太大損失。3、卸載（Offloading）是將未使用的數(shù)據(jù)暫時卸載到CPU或不同的設(shè)備中，然后在需要時再將其讀回。樸素的實現(xiàn)方式會大大降低訓(xùn)練速度，但復(fù)雜的實現(xiàn)方式會預(yù)先獲取數(shù)據(jù)，這樣設(shè)備就不需要再等待了。這個想法的一個具體實現(xiàn)是ZeRO，它將參數(shù)、梯度和優(yōu)化器狀態(tài)分割到所有可用的硬件上，并根據(jù)實際需要再將它們具體化。4、內(nèi)存效率優(yōu)化器（Memory Efficient Optimizer）可以減少優(yōu)化器所維護(hù)的運行狀態(tài)的內(nèi)存占用，如Adafactor。5、壓縮（Compression）也可用于存儲網(wǎng)絡(luò)中的中間結(jié)果。例如，Gist對為后向傳遞而保存的激活進(jìn)行壓縮；DALL-E在同步梯度之前壓縮了梯度。
參考資料：
https://openai.com/blog/techniques-for-training-large-neural-networks/
