本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

什么？單塊GPU也能訓(xùn)練大模型了？

還是20系就能拿下的那種？？？

沒(méi)開(kāi)玩笑，事實(shí)已經(jīng)擺在眼前：

• RTX 2060 6GB普通游戲本能訓(xùn)練15億參數(shù)模型；
• RTX 3090 24GB主機(jī)直接單挑180億參數(shù)大模型；
• Tesla V100 32GB連240億參數(shù)都能拿下。

相比于PyTorch和業(yè)界主流的DeepSpeed方法，提升參數(shù)容量能達(dá)到10多倍。

而且這種方法完全開(kāi)源，只需要幾行代碼就能搞定，修改量也非常少。

這波操作真是直接腰斬大模型訓(xùn)練門(mén)檻啊，老黃豈不是要血虧。

那么，搞出如此大名堂的是何方大佬呢？

它就是國(guó)產(chǎn)開(kāi)源項(xiàng)目Colossal-AI。

自開(kāi)源以來(lái)，曾多次霸榜GitHub熱門(mén)第一。

△開(kāi)源地址：https://github.com/hpcaitech/ColossalAI

主要做的事情就是加速各種大模型訓(xùn)練，GPT-2、GPT-3、ViT、BERT等模型都能搞定。

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

同時(shí)還能省GPU。

比如訓(xùn)練GPT-3時(shí)使用的GPU資源，可以只是英偉達(dá)Megatron-LM的一半。

那么這一回，它又是如何讓單塊GPU訓(xùn)練百億參數(shù)大模型的呢？

我們深扒了一下原理~

高效利用GPU+CPU異構(gòu)內(nèi)存

為什么單張消費(fèi)級(jí)顯卡很難訓(xùn)練AI大模型？

顯存有限，是最大的困難。

當(dāng)今大模型風(fēng)頭正盛、效果又好，誰(shuí)不想上手感受一把？

但動(dòng)不動(dòng)就“CUDA out of memory”，著實(shí)讓人遭不住。

目前，業(yè)界主流方法是微軟DeepSpeed提出的ZeRO （Zero Reduency Optimizer）。

它的主要原理是將模型切分，把模型內(nèi)存平均分配到單個(gè)GPU上。

數(shù)據(jù)并行度越高，GPU上的內(nèi)存消耗越低。

這種方法在CPU和GPU內(nèi)存之間僅使用靜態(tài)劃分模型數(shù)據(jù)，而且內(nèi)存布局針對(duì)不同的訓(xùn)練配置也是恒定的。

由此會(huì)導(dǎo)致兩方面問(wèn)題。

第一，當(dāng)GPU或CPU內(nèi)存不足以滿足相應(yīng)模型數(shù)據(jù)要求時(shí)，即使還有其他設(shè)備上有內(nèi)存可用，系統(tǒng)還是會(huì)崩潰。

第二，細(xì)粒度的張量在不同內(nèi)存空間傳輸時(shí)，通信效率會(huì)很低；當(dāng)可以將模型數(shù)據(jù)提前放置到目標(biāo)計(jì)算設(shè)備上時(shí)，CPU-GPU的通信量又是不必要的。

目前已經(jīng)出現(xiàn)了不少DeepSpeed的魔改版本，提出使用電腦硬盤(pán)來(lái)動(dòng)態(tài)存儲(chǔ)模型，但是硬盤(pán)的讀寫(xiě)速度明顯低于內(nèi)存和顯存，訓(xùn)練速度依舊會(huì)被拖慢。

針對(duì)這些問(wèn)題，Colossal-AI采用的解決思路是高效利用GPU+CPU的異構(gòu)內(nèi)存。

具體來(lái)看，是利用深度學(xué)習(xí)網(wǎng)絡(luò)訓(xùn)練過(guò)程中不斷迭代的特性，按照迭代次數(shù)將整個(gè)訓(xùn)練過(guò)程分為預(yù)熱和正式兩個(gè)階段。

預(yù)熱階段，監(jiān)測(cè)采集到非模型數(shù)據(jù)內(nèi)存信息；

正式階段，根據(jù)采集到的信息，預(yù)留出下一個(gè)算子在計(jì)算設(shè)備上所需的峰值內(nèi)存，移動(dòng)出一些GPU模型張量到CPU內(nèi)存。

大概邏輯如下所示：

這里稍微展開(kāi)說(shuō)明下，模型數(shù)據(jù)由參數(shù)、梯度和優(yōu)化器狀態(tài)組成，它們的足跡和模型結(jié)構(gòu)定義有關(guān)。

非模型數(shù)據(jù)由operator生成的中間張量組成，會(huì)根據(jù)訓(xùn)練任務(wù)的配置（如批次大小）動(dòng)態(tài)變化。

它倆常干的事呢，就是搶GPU顯存。

所以，就需要在GPU顯存不夠時(shí)CPU能來(lái)幫忙，與此同時(shí)還要避免其他情況下內(nèi)存浪費(fèi)。

Colossal-AI高效利用GPU+CPU的異構(gòu)內(nèi)存，就是這樣的邏輯。

而以上過(guò)程中，獲取非模型數(shù)據(jù)的內(nèi)存使用量其實(shí)非常難。

因?yàn)榉悄Ｐ蛿?shù)據(jù)的生存周期并不歸用戶(hù)管理，現(xiàn)有的深度學(xué)習(xí)框架沒(méi)有暴露非模型數(shù)據(jù)的追蹤接口給用戶(hù)。其次，CUDA context等非框架開(kāi)銷(xiāo)也需要統(tǒng)計(jì)。

在這里Colossal-AI的解決思路是，在預(yù)熱階段用采樣的方式，獲得非模型數(shù)據(jù)對(duì)CPU和GPU的內(nèi)存的使用情況。

簡(jiǎn)單來(lái)說(shuō)，這是道加減法運(yùn)算：

非數(shù)據(jù)模型使用 ＝ 兩個(gè)統(tǒng)計(jì)時(shí)刻之間系統(tǒng)最大內(nèi)存使用 — 模型數(shù)據(jù)內(nèi)存使用

已知，模型數(shù)據(jù)內(nèi)存使用可以通過(guò)查詢(xún)管理器得知。

具體來(lái)看就是下面醬嬸的：

所有模型數(shù)據(jù)張量交給內(nèi)存管理器管理，每個(gè)張量標(biāo)記一個(gè)狀態(tài)信息，包括HOLD、COMPUTE、FREE等。

然后，根據(jù)動(dòng)態(tài)查詢(xún)到的內(nèi)存使用情況，不斷動(dòng)態(tài)轉(zhuǎn)換張量狀態(tài)、調(diào)整張量位置，更高效利用GPU顯存和CPU內(nèi)存。

在硬件非常有限的情況下，最大化模型容量和平衡訓(xùn)練速度。這對(duì)于AI普及化、低成本微調(diào)大模型下游任務(wù)等，都具有深遠(yuǎn)意義。

而且最最最關(guān)鍵的是——加內(nèi)存條可比買(mǎi)高端顯卡劃 算 多 了。

前不久，Colossal-AI還成功復(fù)現(xiàn)了谷歌的最新研究成果PaLM (Pathways Language Model)，表現(xiàn)同樣非常奈斯，而微軟DeepSpeed目前還不支持PaLM模型。

Colossal-AI還能做什么？

前面也提到，Colossal-AI能挑戰(zhàn)的任務(wù)非常多，比如加速訓(xùn)練、節(jié)省GPU資源。

那么它是如何做到的呢？

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

目前，基于Colossal-AI的加速方案FastFold，能夠?qū)⒌鞍踪|(zhì)結(jié)構(gòu)預(yù)測(cè)模型AlphaFold的訓(xùn)練時(shí)間，從原本的11天，減少到只需67小時(shí)。

而且總成本更低，在長(zhǎng)序列推理任務(wù)中，也能實(shí)現(xiàn)9~11.6倍的速度提升。

這一方案成功超越谷歌和哥倫比亞大學(xué)的方法。

此外，Colossal-AI還能只用一半GPU數(shù)量訓(xùn)練GPT-3。

相比英偉達(dá)方案，Colossal-AI僅需一半的計(jì)算資源，即可啟動(dòng)訓(xùn)練；若使用相同計(jì)算資源，則能提速11%，可降低GPT-3訓(xùn)練成本超百萬(wàn)美元。

與此同時(shí)，Colossal-AI也非常注重開(kāi)源社區(qū)建設(shè)，提供中文教程、開(kāi)放用戶(hù)社群論壇，根據(jù)大家的需求反饋不斷更新迭代。

比如之前有讀者留言說(shuō)，Colossal-AI要是能在普通消費(fèi)級(jí)顯卡上跑就好了。

這不，幾個(gè)月后，已經(jīng)安排好了~

背后團(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)化器的提出者尤洋。

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

拿過(guò)IPDPS/ICPP最佳論文、ACM/IEEE George Michael HPC Fellowship、福布斯30歲以下精英（亞洲 2021）、IEEE-CS超算杰出新人獎(jiǎng)、UC伯克利EECS Lotfi A. Zadeh優(yōu)秀畢業(yè)生獎(jiǎng)。

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

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

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

團(tuán)隊(duì)的核心成員均來(lái)自美國(guó)加州大學(xué)伯克利分校、哈佛大學(xué)、斯坦福大學(xué)、芝加哥大學(xué)、清華大學(xué)、北京大學(xué)、新加坡國(guó)立大學(xué)、新加坡南洋理工大學(xué)等國(guó)內(nèi)外知名高校；擁有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名廠商工作經(jīng)歷。

公司成立即獲得創(chuàng)新工場(chǎng)、真格基金等多家頂尖VC機(jī)構(gòu)種子輪投資。

潞晨CSO Prof. James Demmel為加州大學(xué)伯克利分校杰出教授、ACM/IEEE Fellow，同時(shí)還是美國(guó)科學(xué)院、工程院、藝術(shù)與科學(xué)院三院院士。

傳送門(mén)：

??https://github.com/hpcaitech/ColossalAI??