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