要了解大模型訓(xùn)練難，我們得先看看從傳統(tǒng)的分布式訓(xùn)練，到大模型的出現(xiàn)逞敷，需要大規(guī)模分布式訓(xùn)練的原因。接著第二點(diǎn)去了解下大規(guī)模訓(xùn)練的挑戰(zhàn)灌侣。

從分布式訓(xùn)練到大規(guī)模訓(xùn)練

常見的訓(xùn)練方式是單機(jī)單卡推捐，也就是一臺服務(wù)器配置1塊AI芯片，這是最簡單的訓(xùn)練方式侧啼。隨著數(shù)據(jù)量的增加牛柒，希望加快模型的訓(xùn)練速度，于是出現(xiàn)了單機(jī)多卡痊乾，多塊AI芯片并行皮壁，以一臺機(jī)器上配置8塊AI芯片為例，把數(shù)據(jù)切分成8份哪审，分別在8塊AI芯片上都跑一次BP算法蛾魄，計(jì)算出梯度，然后所有AI芯片上計(jì)算出的梯度進(jìn)行平均协饲，更新模型參數(shù)畏腕。這樣的話，以前一次BP只能訓(xùn)練1個(gè)batch的數(shù)據(jù)茉稠，現(xiàn)在就是8個(gè)batch描馅。

image

模型和數(shù)據(jù)規(guī)模的增大，意味著訓(xùn)練時(shí)間的增長而线。為了提升模型訓(xùn)練的速度铭污，可以增加計(jì)算資源來縮短訓(xùn)練時(shí)間，于是出現(xiàn)了分布式訓(xùn)練膀篮。所謂分布式訓(xùn)練嘹狞，物理上就是使用多臺機(jī)器，每臺機(jī)器上都有多塊AI芯片誓竿，網(wǎng)絡(luò)模型運(yùn)行在不同機(jī)器的不同AI芯片上磅网，以加快整體訓(xùn)練速度。簡單來理解筷屡，實(shí)質(zhì)就是將單卡的負(fù)載拆到了多卡上涧偷。

image

來看看引入分布式訓(xùn)練的兩個(gè)原因：

1）第一個(gè)是數(shù)據(jù)規(guī)模大簸喂，導(dǎo)致訓(xùn)練時(shí)間很長。在單機(jī)8卡情況下燎潮，對于COCO在115k規(guī)模的數(shù)據(jù)集喻鳄，假設(shè)訓(xùn)練resnet152模型需要40個(gè)小時(shí)。訓(xùn)練Open Image dataset v4圖片數(shù)據(jù)集在1,740k規(guī)模下确封，則相對來說可能需要40天時(shí)間除呵。對于生命寶貴、畢業(yè)延期的煉丹同學(xué)來說爪喘，需要不斷的調(diào)整網(wǎng)絡(luò)結(jié)構(gòu)颜曾、嘗試新的參數(shù)，修改下游任務(wù)等工作腥放，單機(jī)單卡和單機(jī)多卡的訓(xùn)練速度是無法接受的泛啸，于是很有必要使用分布式進(jìn)行訓(xùn)練。

2）第二點(diǎn)就是分布式訓(xùn)練可能帶來精度上的提升秃症。先回憶一下，為什么要用SGD來優(yōu)化模型吕粹，隨機(jī)梯度下降的“隨機(jī)”是指每次從數(shù)據(jù)集里面隨機(jī)抽取一個(gè)小batch數(shù)據(jù)來進(jìn)行計(jì)算真實(shí)值和模型預(yù)測值之間的誤差种柑，然后反向傳播。之所以只選一個(gè)小batch匹耕，一是因?yàn)殡S機(jī)數(shù)據(jù)產(chǎn)生的小batch梯度方向聚请，理論和實(shí)踐證明，一定程度上可以模擬整個(gè)數(shù)據(jù)的梯度下降方向稳其，二是因?yàn)锳I芯片的內(nèi)存有限驶赏。但是在實(shí)際情況下，小batch梯度并不足夠代替整個(gè)數(shù)據(jù)集的梯度方向既鞠，每次小batch在BP算法求解出來的梯度方向煤傍，與整體數(shù)據(jù)集并不完全一致。這樣就會導(dǎo)致優(yōu)化迭代（訓(xùn)練）過程不斷震蕩嘱蛋。使用分布式訓(xùn)練蚯姆，可以使用更大的batch size，避免優(yōu)化過程中的震蕩洒敏。

數(shù)據(jù)并行：在分布式訓(xùn)練場景龄恋，因?yàn)閿?shù)據(jù)規(guī)模增大，可以通過數(shù)據(jù)并行凶伙，可以對輸入數(shù)據(jù)進(jìn)行切分郭毕，每塊AI芯片只需要處理一部分?jǐn)?shù)據(jù)。優(yōu)點(diǎn)是計(jì)算通信可以重疊函荣，通信需求低显押，并行效率高扳肛；缺點(diǎn)是無法運(yùn)行大模型。

image

模型并行：同時(shí)煮落，模型每一層的參數(shù)量增大敞峭、模型的層數(shù)越深，可以通過模型并行修改網(wǎng)絡(luò)層內(nèi)的計(jì)算方式蝉仇，將單層的計(jì)算負(fù)載和內(nèi)存負(fù)載切分到多塊AI芯片上旋讹。缺點(diǎn)是通信和計(jì)算是串行的，對通信的要求高轿衔。

image

流水線并行：另外還可以利用流水并行沉迹，將不同的網(wǎng)絡(luò)層放在不同的AI芯片上運(yùn)行，進(jìn)而進(jìn)一步將計(jì)算負(fù)載和內(nèi)存負(fù)載切分至多塊AI芯片害驹。優(yōu)點(diǎn)是計(jì)算的通信可以重疊鞭呕，通信需求低；但是存在流水線間空閑bubble宛官，需要和重計(jì)算配合使用葫松。

image

通過對數(shù)據(jù)、網(wǎng)絡(luò)模型底洗、以及運(yùn)行的流水線進(jìn)行切分腋么，分布式訓(xùn)練可以有效地減少單卡的資源負(fù)載，一方面提升了訓(xùn)練任務(wù)的計(jì)算和內(nèi)存吞吐量亥揖，另一方面使得原來單機(jī)單卡無法訓(xùn)練的任務(wù)珊擂，變成可能。

現(xiàn)在一切看上去都是那么的美好费变，有分布式訓(xùn)練就夠了摧扇。理論上，AI芯片數(shù)量越多挚歧，模型訓(xùn)練越快扛稽。但是，隨著訓(xùn)練數(shù)據(jù)集規(guī)模的進(jìn)一步增長昼激，數(shù)據(jù)并行就會出現(xiàn)局限性：當(dāng)訓(xùn)練資源擴(kuò)大到一定規(guī)模時(shí)庇绽，由于通信瓶頸的存在，增加計(jì)算資源的邊際效應(yīng)橙困，就會越來越明顯瞧掺，甚至增加資源也沒辦法進(jìn)行加速，這個(gè)比例稱為加速比凡傅。假設(shè)單設(shè)備的吞吐量為T辟狈，有n個(gè)設(shè)備系統(tǒng)的吞吐量應(yīng)為nT，設(shè)這個(gè)系統(tǒng)實(shí)際達(dá)到的吞吐量為T_n。最終的加速比(scale factor)為：

scale factor=T_n/nT

image

理論上哼转，我們希望無論有設(shè)備數(shù)n為多少明未，加速比越接近1越好。

另外壹蔓，隨著網(wǎng)絡(luò)模型的不斷增大趟妥，會出現(xiàn)模型需要的內(nèi)存急劇膨脹，算子的增加也會使得單卡就算進(jìn)行模型切分和流水線切分佣蓉，也難以在合理的時(shí)間內(nèi)完成一個(gè)step的訓(xùn)練披摄，最終導(dǎo)致分布式訓(xùn)練不再適用于這類型的任務(wù)。這也是大模型的出現(xiàn)勇凭，對系統(tǒng)的挑戰(zhàn)疚膊。

這個(gè)時(shí)候，面對大模型虾标，我們需要引入大規(guī)模訓(xùn)練技術(shù)寓盗，在解決內(nèi)存的同時(shí)，也不會被計(jì)算資源給限制住璧函，讓算法開發(fā)人員可以方便進(jìn)行高效的分布式調(diào)試調(diào)優(yōu)和編寫代碼傀蚌。愉快地訓(xùn)練大模型。

大規(guī)模訓(xùn)練的挑戰(zhàn)

相比普通的分布式訓(xùn)練蘸吓，大規(guī)模訓(xùn)練在技術(shù)上喳张，需要考慮的問題更加復(fù)雜。

首先美澳，面對單卡無法裝載的大模型，如何利用多卡來突破內(nèi)存限制的瓶頸是個(gè)問題摸航；其次制跟，大規(guī)模訓(xùn)練會用到大量的計(jì)算資源，大量計(jì)算資源間如何通信酱虎、協(xié)作是另一個(gè)難題雨膨；最后，如何平衡各類層出不窮的大規(guī)模訓(xùn)練技術(shù)读串，使得眾多技術(shù)形成一個(gè)完整高效的訓(xùn)練方案聊记，更是一大學(xué)問。

下面恢暖，將大規(guī)模訓(xùn)練技術(shù)面臨的挑戰(zhàn)分為四個(gè)部分：內(nèi)存排监、通訊、計(jì)算和調(diào)優(yōu)杰捂。

image

內(nèi)存墻

模型訓(xùn)練無可避免的問題就是內(nèi)存墻舆床。

下面舉兩個(gè)例子，在進(jìn)行ResNet152對120G的ImageNet數(shù)據(jù)進(jìn)行訓(xùn)練時(shí)，在使用Batch Size=128挨队，Input Image Size=512x512的前提下谷暮，訓(xùn)練過程占用了16G的內(nèi)存，在GPU一般具有16G內(nèi)存的情況下盛垦，沒辦法把模型放在一個(gè)一塊AI芯片上湿弦。想進(jìn)一步加大每次執(zhí)行數(shù)據(jù)的Batch size，模型訓(xùn)練的內(nèi)存占用也會隨之增長腾夯，最后高于AI芯片的內(nèi)存容量颊埃，觸碰到了內(nèi)存墻，模型無法訓(xùn)練俯在。

以2000億參數(shù)的鵬程盤古大模型為例竟秫，2000億的參數(shù)內(nèi)存占用就消耗了754GB，訓(xùn)練過程因?yàn)闀袡?quán)重跷乐、激活肥败、優(yōu)化器狀態(tài)、再加上自動(dòng)微分所產(chǎn)生臨時(shí)變量愕提，需要3500GB的內(nèi)存馒稍，一個(gè)大模型訓(xùn)練就需要100多塊具有32G內(nèi)存的AI芯片。

要理解內(nèi)存墻出現(xiàn)的本質(zhì)浅侨，就需要了解內(nèi)存增長的本質(zhì)纽谒。在ResNet50的1輪迭代，打開看看內(nèi)存占用變化：

1）網(wǎng)絡(luò)模型開始計(jì)算的時(shí)候如输，內(nèi)存占用不斷增加鼓黔，直到達(dá)到峰值1.2GB。

2）峰值過后內(nèi)存開始逐漸釋放不见，內(nèi)存占用慢慢降到320M澳化。

3）1個(gè)step計(jì)算結(jié)束后，仍有一部分內(nèi)存駐留稳吮，內(nèi)存保持在320M缎谷。

image

可以發(fā)現(xiàn)，模型訓(xùn)練對內(nèi)存的占用可以分為兩部分：

1）一部分是模型自身的權(quán)重參數(shù)灶似、優(yōu)化器狀態(tài)信息列林，由于是比較固定的所以稱為靜態(tài)參數(shù)。

2）另一部分是模型在前向計(jì)算和反向傳播的時(shí)候酪惭，會產(chǎn)生諸如前向輸出Forward Output希痴、梯度輸出Output Gradient、算子計(jì)算時(shí)候的臨時(shí)變量等撞蚕，這部分內(nèi)存會在反向傳播時(shí)逐漸釋放掉润梯，因此被稱為動(dòng)態(tài)參數(shù)。

從這里面可以發(fā)現(xiàn)，在計(jì)算過程的內(nèi)存峰值纺铭，是否遇到內(nèi)存墻寇钉，主要是由動(dòng)態(tài)內(nèi)存決定。即使絞盡腦汁舶赔，降低靜態(tài)內(nèi)存意義并不是很大扫倡，關(guān)鍵在于降低內(nèi)存占用的峰值。

靜態(tài)內(nèi)存

以ResNet50為例竟纳，靜態(tài)內(nèi)存占用大約1.2G：

ResNet50 = 250M(權(quán)重參數(shù)) + 900M(優(yōu)化器參數(shù)和動(dòng)量) ~ 1.2G

在計(jì)算過程中撵溃，神經(jīng)網(wǎng)絡(luò)模型每一層的卷積或者全連接計(jì)算，都會把權(quán)重W_m長期保存下來锥累，用作網(wǎng)絡(luò)的權(quán)重參數(shù)更新缘挑。另外針對諸如ADAM的優(yōu)化器，會存儲優(yōu)化器的動(dòng)量等信息桶略，用于優(yōu)化器計(jì)算语淘。

一塊AI芯片有16G，最大能塞滿20+億參數(shù)的模型际歼，但是這時(shí)候已經(jīng)沒有額外空間惶翻，留給動(dòng)態(tài)內(nèi)存進(jìn)行分配啦。所以說降低動(dòng)態(tài)內(nèi)存的峰值才是人間正道鹅心。

image

動(dòng)態(tài)內(nèi)存

靜態(tài)內(nèi)存比較好理解吕粗，下面主要看看動(dòng)態(tài)內(nèi)存。動(dòng)態(tài)內(nèi)存一般指的在自動(dòng)微分Autogard過程旭愧，產(chǎn)生的中間變量颅筋。

神經(jīng)網(wǎng)絡(luò)層數(shù)非常多，在計(jì)算過程中输枯，神經(jīng)網(wǎng)絡(luò)模型每一層的前向輸出Forward Output都需要保存下來給反向傳播時(shí)候使用垃沦；而在神經(jīng)網(wǎng)絡(luò)的反向傳播計(jì)算過程中，需要記錄下權(quán)重梯度Weight Gradient(W_g)用押，給優(yōu)化器使用，在后續(xù)的step中更新梯度的信息靶剑；另外還有梯度輸出Output Gradient同樣在反向傳播的過程中蜻拨，繼續(xù)在反向傳播的過程中，作為上一層網(wǎng)絡(luò)模型的輸入桩引。

當(dāng)然缎讼，這里面還會有一些正向和反向的算子中間變量需要消耗內(nèi)存的，由此坑匠，上述的O_f血崭、O_g、W_g、OP累積下來,會造成巨大的內(nèi)存開銷夹纫。從ResNet50的例子中咽瓷，我們可以看出，動(dòng)態(tài)內(nèi)存是靜態(tài)內(nèi)存的X倍(舰讹，靜態(tài)內(nèi)存占用了1.2G內(nèi)存茅姜，而動(dòng)態(tài)內(nèi)存則占用了25G內(nèi)存，比例非常不平衡)月匣。

image

靜態(tài)內(nèi)存和動(dòng)態(tài)內(nèi)存都可能造成內(nèi)存墻的問題钻洒。相互獨(dú)立但又相互制約，任意一側(cè)的增大都會導(dǎo)致留給另一側(cè)的顯存空間變小锄开，所以單單對一側(cè)做優(yōu)化是不夠的素标，必須同時(shí)優(yōu)化靜態(tài)和動(dòng)態(tài)內(nèi)存。

為了能夠?qū)⒋竽Ｐ瓦\(yùn)行起來萍悴，需要使用模型并行头遭、流水并行等技術(shù)，但是這些技術(shù)又會降低AI芯片的計(jì)算吞吐量退腥。

通訊墻

大模型通過模型并行任岸、流水線并行切分到AI集群后，通訊便成了主要的性能瓶頸狡刘。

從模型切分的角度來考慮享潜，大模型再怎么切分其仍然是一個(gè)模型，因而模型切分到不同的機(jī)器后嗅蔬，仍然需要通信來對分布在不同機(jī)器的參數(shù)進(jìn)行總體聚合剑按。

在通訊方式的角度來看，參數(shù)聚合(All Reduce)就會對通訊提出很多需求澜术，例如使用同步的更新策略艺蝴，還是使用異步的更新策略祈秕，如何對模型局部變量進(jìn)行更新等乡小。

另外，在網(wǎng)絡(luò)通訊的角度耐量，由于專用的AI加速芯片中內(nèi)存與計(jì)算單元ALU非常進(jìn)盒延，使得片內(nèi)帶寬很大缩擂，計(jì)算速度很快，但是在集群中的網(wǎng)絡(luò)傳輸速度遠(yuǎn)遠(yuǎn)不能匹配專用的AI加速芯片運(yùn)算速率添寺。

直接用更大的帶寬不能解決這些問題嗎胯盯？

從圖中，我們可以看到计露，隨著網(wǎng)絡(luò)帶寬從1Gbps到100Gbps博脑，利用率從接近100%下降到40%左右憎乙。隨著網(wǎng)絡(luò)帶寬的增加，帶寬的利用率將會越來越低叉趣，高帶寬所帶來的收益會遇到瓶頸泞边。

image

隨著機(jī)器規(guī)模的擴(kuò)大，基于同步的All Reduce通訊聚合方式君账，會因?yàn)榇罅康腁I芯片和服務(wù)器之間頻繁進(jìn)行同步繁堡，出現(xiàn)水桶效應(yīng)，也就是最慢的一路通訊乡数，將會決定整個(gè)AI集群的通訊的高度椭蹄。

如果采用目前比較流行的Ring-AllReduce的通信聚合方式，當(dāng)通訊的環(huán)越大净赴，通訊的延長將會不斷地被擴(kuò)大绳矩。另外網(wǎng)絡(luò)協(xié)議的多次握手的方式，諸如此類的開銷會導(dǎo)致訓(xùn)練無法有效利用帶寬玖翅。

總體而言翼馆，通訊過程中，需要綜合考慮數(shù)據(jù)參數(shù)量金度、計(jì)算量应媚、計(jì)算類型、數(shù)據(jù)樣本量猜极、集群帶寬拓?fù)浜屯ㄓ嵅呗缘炔煌囊蛩刂薪拍茉O(shè)計(jì)出性能較優(yōu)的切分策略，最大化利用通訊效率跟伏，提高通訊比丢胚。

性能墻

性能墻呢主要是指計(jì)算資源利用率的問題。隨著大模型的提出受扳，對算力需求更加迫切携龟，理論上在4K的集群上每塊卡快1分鐘，總體就快了68個(gè)小時(shí)勘高。大模型會增加對算力的需求峡蟋，但是隨著大模型引入各項(xiàng)分布式并行技術(shù)的同時(shí)，會降低計(jì)算資源的利用率华望。

計(jì)算效率在MMlab的貼總結(jié)得非常好层亿，我們同樣從底至上分為算子層(Operator Level)、圖層(Graph Level)和任務(wù)層(Task Level)三個(gè)方面立美。

算子層(Operator Level)

在Operator Level的問題大模型和普通模型遇到的問題都是類似，主要是指對算子級別的優(yōu)化方灾。

1）其中比較突出的問題是模型小算子過多建蹄，可以通過算子融合進(jìn)行優(yōu)化碌更，例如對Conv+Bn+Act三個(gè)算子進(jìn)行融合成一個(gè)大算子。但是還有很多算子融合策略無法窮舉洞慎，如何通過計(jì)算編譯技術(shù)痛单，對小算子進(jìn)行融合。

2）第二點(diǎn)就是算子實(shí)現(xiàn)不夠高效劲腿，類似于卷積CONV算子針對2x2和3x3 Kernel大小的使用Winograd算法代替旭绒。可是這種算法是人為提出了焦人，更多的高效Kernel算子挥吵，如何利用編譯技術(shù)，去尋找最好的多線程并行計(jì)算花椭。

3）最后一點(diǎn)就是內(nèi)存局部性差忽匈，對算子和內(nèi)存的開銷進(jìn)行分析，可以對計(jì)算時(shí)算子輸出有相同的shape進(jìn)行復(fù)用矿辽。

圖層(Graph Level)

Graph Level指如何對計(jì)算圖進(jìn)行優(yōu)化丹允，進(jìn)而加速大規(guī)模訓(xùn)練。

一般而言袋倔，如果沒有遇到控制流和大模型雕蔽，一個(gè)模型只有一個(gè)圖。但是當(dāng)遇到大模型的時(shí)候宾娜，由于一塊AI芯片的內(nèi)存沒辦法滿足大模型的內(nèi)存需求批狐，因此需要對算子和模型進(jìn)行切分。這時(shí)候涉及兩方面的問題碳默。

1）如何搜索和切分處計(jì)算效率更高的子圖贾陷，分配到不同的機(jī)器上進(jìn)行計(jì)算。

2）數(shù)據(jù)在通訊和內(nèi)存之間如何增加overlay重疊部分嘱根，提高單卡整體的計(jì)算率髓废。

任務(wù)層Task Level

大集群大模型的分布式混合了各種并行訓(xùn)練策略，在任務(wù)層級该抒，性能的瓶頸從計(jì)算轉(zhuǎn)移到了通信慌洪，如何降低通信量，縮減通信域規(guī)模凑保，盡可能把通信時(shí)延隱藏在計(jì)算中冈爹，是大規(guī)模訓(xùn)練的核心關(guān)注點(diǎn)。

大規(guī)模訓(xùn)練技術(shù)中欧引，不僅要求AI芯片的計(jì)算性能足夠強(qiáng)悍频伤，同時(shí)也依賴于AI框架的大規(guī)模分布式訓(xùn)練的運(yùn)行和調(diào)度效率，以及在分布式并行等各種優(yōu)化手段的權(quán)衡芝此。

調(diào)優(yōu)墻

在數(shù)千節(jié)點(diǎn)的集群上進(jìn)行模型開發(fā)憋肖，聽到就頭皮發(fā)麻因痛。我平時(shí)把EffectNet網(wǎng)絡(luò)模型魔改，還不一定能夠收斂岸更，調(diào)一次參數(shù)再訓(xùn)練一次鸵膏，單機(jī)多卡一個(gè)星期就過去了。那要是大模型訓(xùn)練了1個(gè)月遇到Loss跑飛了怎么辦怎炊？

所以在數(shù)千節(jié)點(diǎn)的集群上谭企，需要考慮到提升算法工程師分布式調(diào)試調(diào)優(yōu)的效率，另外還要考慮降低工程師對大模型進(jìn)行并行切分的難度评肆。除了對人的考慮债查，還要對硬件集群的管理，需要保證計(jì)算的正確性糟港、性能攀操、可用性。要是有一臺機(jī)器壞了秸抚，如何快速恢復(fù)訓(xùn)練中的參數(shù)速和。

總結(jié)一句話，大模型訓(xùn)練考驗(yàn)的是算法剥汤、數(shù)據(jù)颠放、框架、資源調(diào)度等全棧和全流程的綜合能力吭敢。希望MindSpore能夠在大模型引領(lǐng)業(yè)界碰凶。