2天訓(xùn)練出15億參數(shù)大模型，國(guó)產(chǎn)開(kāi)源項(xiàng)目力克英偉達(dá)Megatron-LM，來(lái)自LAMB作者團(tuán)隊(duì)

魚(yú)羊 明敏 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

當(dāng)今AI之勢(shì)，影響縱深發(fā)展的矛盾是什么？

一方面，大模型風(fēng)頭正勁，效果驚艷，人人都想試試。但另一方面，硬件基礎(chǔ)上動(dòng)不動(dòng)就是上萬(wàn)張GPU的大規(guī)模集群在日夜燃燒，鈔能力勸退。

所以如果告訴你，現(xiàn)在只用一半數(shù)量的GPU，也能完成同樣的GPT-3訓(xùn)練呢？

你會(huì)覺(jué)得關(guān)鍵鑰匙是什么？

不賣(mài)關(guān)子了。實(shí)現(xiàn)如此提升的，是一個(gè)名為Colossal-AI的GitHub開(kāi)源項(xiàng)目。

而且該項(xiàng)目開(kāi)源不久，就迅速登上了Python方向的熱榜世界第一。

↑GitHub地址：https://github.com/hpcaitech/ColossalAI

不僅能加速GPT-3，對(duì)于GPT-2、ViT、BERT等多種模型，Colossal-AI的表現(xiàn)也都非常nice：

比如半小時(shí)左右就能預(yù)訓(xùn)練一遍ViT-Base/32，2天能訓(xùn)完15億參數(shù)GPT模型、5天可訓(xùn)完83億參數(shù)GPT模型。

與業(yè)內(nèi)主流的AI并行系統(tǒng)——英偉達(dá)Megatron-LM相比，在同樣使用512塊GPU訓(xùn)練GPT-2模型時(shí)，Colossal-AI的加速比是其2倍。而在訓(xùn)練GPT-3時(shí)，更是可以節(jié)省近千萬(wàn)元的訓(xùn)練費(fèi)用。

此外在訓(xùn)練GPT-2時(shí)，顯存消耗甚至能控制在Megatron-LM的十分之一以下。

Colossal-AI究竟是如何做到的？

老規(guī)矩，我們從論文扒起。

高效6維并行方法

簡(jiǎn)單來(lái)說(shuō)，Colossal-AI就是一個(gè)整合了多種并行方法的系統(tǒng)，提供的功能包括多維并行、大規(guī)模優(yōu)化器、自適應(yīng)任務(wù)調(diào)度、消除冗余內(nèi)存等。

首先來(lái)看多維并行。

所謂“多維”是指，目前主流的分布式并行方案往往使用多種并行方法。

比如英偉達(dá)的Megatron-LM使用了3種方法：數(shù)據(jù)并行、流水并行和張量并行。因此這種模式也被稱(chēng)為三維并行。微軟的DeepSpeed調(diào)用Megatron-LM作為并行基礎(chǔ)。

而Colossal-AI能將系統(tǒng)的并行維度，一下子拉升到6維——

在兼容數(shù)據(jù)并行、流水并行的基礎(chǔ)上，基于該項(xiàng)目團(tuán)隊(duì)自研的2維/2.5維/3維張量并行方法，以及序列并行實(shí)現(xiàn)。

其中，高維張量并行正是Colossal-AI提升大模型顯存利用率和通信效率的關(guān)鍵所在。

其實(shí)張量并行并不新奇，只是過(guò)去我們常見(jiàn)的張量并行更多都是基于一維的。

它的原理是將模型層內(nèi)的權(quán)重參數(shù)按行或列切分到不同的處理器上，利用分塊矩陣乘法，將一個(gè)運(yùn)算分布到多個(gè)處理器上同時(shí)進(jìn)行。

比如英偉達(dá)的Megatron-LM就是一個(gè)典型的例子。

但這種并行方式存在一定弊端。

比如，每個(gè)處理器仍需要存儲(chǔ)整個(gè)中間激活，使得在處理大模型時(shí)會(huì)浪費(fèi)大量顯存空間。

另一方面，這種單線方法還會(huì)導(dǎo)致每個(gè)處理器都需要與其他所有處理器進(jìn)行通信。

這意味著假設(shè)有100個(gè)GPU的話，每個(gè)GPU都需要與其他99個(gè)GPU通信，每次計(jì)算需要通信的次數(shù)就高達(dá)9900次。

但如果將張量并行的維度擴(kuò)展到2維，單次計(jì)算量能立刻下降一個(gè)量級(jí)。

因?yàn)槊總€(gè)GPU只需與自己同行或同列的GPU通信即可。

同樣還是100個(gè)GPU的情況，每個(gè)GPU需要通信的GPU個(gè)數(shù)就能降到9個(gè)，單次計(jì)算僅需900次。

實(shí)際上在此基礎(chǔ)上，Colossal-AI還包含2.5維、3維張量并行方法，可以進(jìn)一步降低傳輸成本。

相較于2維并行方法，2.5維并行方法可提升1.45倍效率，3維方法可提升1.57倍。

針對(duì)大圖片、視頻、長(zhǎng)文本、長(zhǎng)時(shí)間醫(yī)療監(jiān)控等數(shù)據(jù)，Colossal-AI還使用了序列并行的方法，這種方法能突破原有機(jī)器能力限制，直接處理長(zhǎng)序列數(shù)據(jù)。

值得一提的是，Colossal-AI的API接口是可以定制的，這使得它可以便捷添加新的并行維度。

其次，大規(guī)模優(yōu)化器也是Colossal-AI的亮點(diǎn)。

上面我們也提到了，在分布式并行系統(tǒng)中會(huì)使用多種并行方法，數(shù)據(jù)并行則是另一種常見(jiàn)方法。

這種方法的原理不難理解，就是把訓(xùn)練數(shù)據(jù)劃分成若干份，讓不同的機(jī)器運(yùn)算不同的數(shù)據(jù)，然后通過(guò)一個(gè)參數(shù)服務(wù)器 （Paremeter Server）收集目標(biāo)數(shù)據(jù)。

由此可以大幅提升AI模型訓(xùn)練過(guò)程中的批量大小，加速訓(xùn)練過(guò)程。

不過(guò)大批量訓(xùn)練有個(gè)“通病”，就是會(huì)產(chǎn)生泛化誤差 （Generalization Gap），導(dǎo)致網(wǎng)絡(luò)泛化能力下降，進(jìn)而導(dǎo)致AI模型準(zhǔn)確度下降。

所以，Colossal-AI在系統(tǒng)中使用了自研的LAMB、LARS等大規(guī)模優(yōu)化器。在保證訓(xùn)練精度的情況下，還將批大小從512擴(kuò)展到65536。

其中，LARS優(yōu)化器是通過(guò)逐層調(diào)整學(xué)習(xí)率，來(lái)減少因?yàn)閷W(xué)習(xí)率導(dǎo)致的無(wú)法收斂情況。

LAMB優(yōu)化器則是在LARS的基礎(chǔ)上，將逐層調(diào)整學(xué)習(xí)率的思想應(yīng)用到自適應(yīng)梯度上。

由此，LAMB能夠很好解決此前LARS在BERT訓(xùn)練中存在差異的問(wèn)題，最大批量達(dá)到了64K。

此前，LAMB優(yōu)化器曾成功將預(yù)訓(xùn)練一遍BERT的時(shí)間，從原本的三天三夜縮短到一個(gè)多小時(shí)。

第三方面，Colossal-AI使用自適應(yīng)可擴(kuò)展調(diào)度器來(lái)高效處理任務(wù)。

與現(xiàn)有常見(jiàn)的任務(wù)調(diào)度器不同，Colossal-AI不是靜態(tài)地通過(guò)GPU個(gè)數(shù)來(lái)判斷任務(wù)規(guī)模，而是根據(jù)批大小來(lái)動(dòng)態(tài)、自動(dòng)管理每個(gè)任務(wù).

通過(guò)演化算法，該任務(wù)調(diào)度器還能不斷優(yōu)化調(diào)度決策，更大程度提升GPU利用率。

評(píng)估結(jié)果表明，與當(dāng)前最先進(jìn)的方法相比，該方法在平均JCT （job completion time）上能夠縮短45.6%的時(shí)間，優(yōu)于現(xiàn)有的深度學(xué)習(xí)任務(wù)調(diào)度算法。

此外，這種自適應(yīng)可擴(kuò)展調(diào)度器還能通過(guò)NCCL網(wǎng)絡(luò)通信實(shí)現(xiàn)高效的任務(wù)遷移。

最后，消除冗余內(nèi)存也是加速AI訓(xùn)練的一種解決思路。

在這方面，Colossal-AI使用了zero redundancy optimizer技術(shù)（簡(jiǎn)稱(chēng)ZeRO）。

這種方法主要通過(guò)切分優(yōu)化器狀態(tài)、梯度、模型參數(shù)，使GPU僅保存當(dāng)前計(jì)算所需的部分，從而來(lái)消除數(shù)據(jù)并行、模型并行中存在的內(nèi)存冗余。

尤其是在部署模型推理時(shí)，通過(guò)zero offload可以將模型卸載到CPU內(nèi)存或硬盤(pán)，僅使用少量GPU資源，即可實(shí)現(xiàn)低成本部署前沿AI大模型。

綜上不難看出，在技術(shù)層面Colossal-AI的加速效果非常明顯。

而在應(yīng)用層面，Colossal-AI的設(shè)計(jì)也顧及了能耗問(wèn)題和易用性兩個(gè)維度。

考慮到數(shù)據(jù)移動(dòng)會(huì)是能耗的主要來(lái)源，Colossal-AI在不增加計(jì)算量的情況下盡可能減少數(shù)據(jù)移動(dòng)量，以此來(lái)降低能耗。

另一方面，作為一個(gè)開(kāi)源給所有人使用的系統(tǒng)，Colossal-AI的使用門(mén)檻不高，即便是沒(méi)有學(xué)習(xí)過(guò)分布式系統(tǒng)的人也能上手操作。

同時(shí)，只需要極少量的代碼改動(dòng)，Colossal-AI就能將已有的單機(jī)代碼快速擴(kuò)展到并行計(jì)算集群上。

最新實(shí)驗(yàn)結(jié)果釋出

Talk is cheap，效果如何，還是得把實(shí)驗(yàn)結(jié)果展開(kāi)來(lái)看。

Colossal-AI近日釋出的最新實(shí)驗(yàn)結(jié)果表明，這一大規(guī)模AI訓(xùn)練系統(tǒng)具有通用性，在GPT-3、GPT-2、ViT、BERT等流行模型上均有亮眼的加速表現(xiàn)。

注：以下GPU均指英偉達(dá)A100。

GPT-3訓(xùn)練速度提高10.7%

英偉達(dá)的Megatron-LM在加速訓(xùn)練GPT-3時(shí)，至少需要128塊GPU才能啟動(dòng)；而從下表可以看出，使用相同的計(jì)算資源，Colossal-AI可以將每次迭代花費(fèi)的時(shí)間從43.1秒降至38.5秒。

這也就意味著，Colossal-AI可以將GPT-3的訓(xùn)練速度進(jìn)一步提高10.7%。

站在工程的角度，考慮到訓(xùn)練這樣的大模型往往需要投入數(shù)百萬(wàn)美元，這一提升比例帶來(lái)的收益不言而喻。

另外，通過(guò)系統(tǒng)優(yōu)化，Colossal-AI還能在訓(xùn)練速度損失不大（43.1→48.5）的前提下，將GPU數(shù)量從128塊減少到96塊，大幅降低訓(xùn)練成本。

而進(jìn)一步啟用ZeRO3（零冗余優(yōu)化器）后，所需GPU數(shù)量甚至能減少一半——至64塊。

2天內(nèi)可完成GPT-2訓(xùn)練

在GPT-2的加速訓(xùn)練結(jié)果中，可以看到，無(wú)論是在4、16還是64塊GPU的情況下，與Megatron-LM相比，Colossal-AI占用的顯存都顯著減少。

也就是說(shuō)，利用Colossal-AI，工程師們可以在采用同等數(shù)量GPU的前提下，訓(xùn)練規(guī)模更大的模型，或設(shè)置更大的批量大小來(lái)加速訓(xùn)練。

從下表結(jié)果中還可以看出，隨著批量大小的增加，Colossal-AI的資源利用率會(huì)進(jìn)一步提高，達(dá)到Megatron-LM速度的2倍。

研發(fā)團(tuán)隊(duì)在256塊GPU上進(jìn)行了實(shí)驗(yàn)，最終用時(shí)82.8個(gè)小時(shí)完成了15億參數(shù)版GPT-2的訓(xùn)練。

據(jù)此預(yù)估，后續(xù)在512塊GPU上進(jìn)行GPT-2預(yù)訓(xùn)練，Colossal-AI能將訓(xùn)練時(shí)間加速到45小時(shí)。

充分兼容多種并行模式

在BERT上進(jìn)行的實(shí)驗(yàn)，則體現(xiàn)了Colossal-AI作為世界上并行維度最多的AI訓(xùn)練系統(tǒng)的優(yōu)勢(shì)。

與Megatron-LM相比，Colossal-AI序列并行方法只需要更少的顯存，就能夠利用更大的批量大小來(lái)加速訓(xùn)練。同時(shí)，還允許開(kāi)發(fā)者使用更長(zhǎng)的序列數(shù)據(jù)。

Colossal-AI的序列并行方法還與流水并行方法兼容。當(dāng)開(kāi)發(fā)者同時(shí)使用序列并行和流水并行時(shí)，可以進(jìn)一步節(jié)省訓(xùn)練大模型的時(shí)間。

另外，在近期的學(xué)術(shù)熱點(diǎn)ViT模型上，Colossal-AI也展現(xiàn)了高維張量并行方法的優(yōu)勢(shì)。

在使用64張GPU的情況下，Colossal-AI采用2/2.5維方式進(jìn)行張量并行，充分利用更大的批量大小，達(dá)到了更快的處理速度。

背后團(tuán)隊(duì)：LAMB優(yōu)化器作者尤洋領(lǐng)銜

看到這里，是不是覺(jué)得Colossal-AI確實(shí)值得標(biāo)星關(guān)注一發(fā)？

實(shí)際上，這一國(guó)產(chǎn)項(xiàng)目背后的研發(fā)團(tuán)隊(duì)來(lái)頭不小。

領(lǐng)銜者，正是LAMB優(yōu)化器的提出者尤洋。

在谷歌實(shí)習(xí)期間，正是憑借LAMB，尤洋曾打破BERT預(yù)訓(xùn)練世界紀(jì)錄。

據(jù)英偉達(dá)官方GitHub顯示，LAMB比Adam優(yōu)化器快出整整72倍。微軟的DeepSpeed也采用了LAMB方法。

說(shuō)回到尤洋本人，他曾以第一名的成績(jī)保送清華計(jì)算機(jī)系碩士研究生，后赴加州大學(xué)伯克利分校攻讀CS博士學(xué)位。

2020年博士畢業(yè)后，他加入新加坡國(guó)立大學(xué)計(jì)算機(jī)系，并于2021年1月成為校長(zhǎng)青年教授（Presidential Young Professor）。

同樣是在2021年，他還獲得了IEEE-CS超算杰出新人獎(jiǎng)。該獎(jiǎng)項(xiàng)每年在全球范圍內(nèi)表彰不超過(guò)3人，僅授予在博士畢業(yè)5年之內(nèi)，已在高性能計(jì)算領(lǐng)域做出有影響力的卓越貢獻(xiàn)，并且可以為高性能計(jì)算的發(fā)展做出長(zhǎng)期貢獻(xiàn)的優(yōu)秀青年學(xué)者。

與此同時(shí)，尤洋回國(guó)創(chuàng)辦潞晨科技——一家主營(yíng)業(yè)務(wù)為分布式軟件系統(tǒng)、大規(guī)模人工智能平臺(tái)以及企業(yè)級(jí)云計(jì)算解決方案的AI初創(chuàng)公司。

其核心團(tuán)隊(duì)成員來(lái)自加州大學(xué)伯克利分校、斯坦福大學(xué)、清華大學(xué)、北京大學(xué)、新加坡國(guó)立大學(xué)、新加坡南洋理工大學(xué)等國(guó)內(nèi)外知名高校，在高性能計(jì)算、人工智能、分布式系統(tǒng)方面有十余年的技術(shù)積累，并已在國(guó)際頂級(jí)學(xué)術(shù)刊物/會(huì)議上發(fā)表論文30余篇。

目前，潞晨科技已拿下創(chuàng)新工場(chǎng)和真格基金合投的超千萬(wàn)元種子輪融資。

傳送門(mén)

有關(guān)Colossal-AI，今天就先介紹到這里。

最后，附上傳送門(mén)，感興趣的小伙伴，自行取用~

GitHub地址：https://github.com/hpcaitech/ColossalAI

參考鏈接：
https://medium.com/@
hpcaitech/efficient-and-easy-training-of-large-ai-models-introducing-colossal-ai-ab571176d3ed