2天訓練出15億引數大模型，國產開源專案力克英偉達Megatron-LM，來自LAMB作者團隊

魚羊 明敏 發自 凹非寺
量子位 | 公眾號 QbitAI

當今AI之勢，影響縱深發展的矛盾是什麼？

一方面，大模型風頭正勁，效果驚豔，人人都想試試。但另一方面，硬體基礎上動不動就是上萬張GPU的大規模叢集在日夜燃燒，鈔能力勸退。

所以如果告訴你，現在只用一半數量的GPU，也能完成同樣的GPT-3訓練呢？

你會覺得關鍵鑰匙是什麼？

不賣關子了。實現如此提升的，是一個名為Colossal-AI的GitHub開源專案。

而且該專案開源不久，就迅速登上了Python方向的熱榜世界第一。

↑GitHub地址：

不僅能加速GPT-3，對於GPT-2、ViT、BERT等多種模型，Colossal-AI的表現也都非常nice：

比如半小時左右就能預訓練一遍ViT-Base/32，2天能訓完15億引數GPT模型、5天可訓完83億引數GPT模型。

與業內主流的AI並行系統——英偉達Megatron-LM相比，在同樣使用512塊GPU訓練GPT-2模型時，Colossal-AI的加速比是其2倍。而在訓練GPT-3時，更是可以節省近千萬元的訓練費用。

此外在訓練GPT-2時，視訊記憶體消耗甚至能控制在Megatron-LM的十分之一以下。

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

老規矩，我們從論文扒起。

高效6維並行方法

簡單來說，Colossal-AI就是一個整合了多種並行方法的系統，提供的功能包括多維並行、大規模優化器、自適應任務排程、消除冗餘記憶體等。

首先來看多維並行。

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

比如英偉達的Megatron-LM使用了3種方法：資料並行、流水並行和張量並行。因此這種模式也被稱為三維並行。微軟的DeepSpeed呼叫Megatron-LM作為並行基礎。

而Colossal-AI能將系統的並行維度，一下子拉昇到6維——

在相容資料並行、流水並行的基礎上，基於該專案團隊自研的2維/2.5維/3維張量並行方法，以及序列並行實現。

其中，高維張量並行正是Colossal-AI提升大模型視訊記憶體利用率和通訊效率的關鍵所在。

其實張量並行並不新奇，只是過去我們常見的張量並行更多都是基於一維的。

它的原理是將模型層內的權重引數按行或列切分到不同的處理器上，利用分塊矩陣乘法，將一個運算分佈到多個處理器上同時進行。

比如英偉達的Megatron-LM就是一個典型的例子。

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

比如，每個處理器仍需要儲存整個中間啟用，使得在處理大模型時會浪費大量視訊記憶體空間。

另一方面，這種單線方法還會導致每個處理器都需要與其他所有處理器進行通訊。

這意味著假設有100個GPU的話，每個GPU都需要與其他99個GPU通訊，每次計算需要通訊的次數就高達9900次。

但如果將張量並行的維度擴充套件到2維，單次計算量能立刻下降一個量級。

因為每個GPU只需與自己同行或同列的GPU通訊即可。

同樣還是100個GPU的情況，每個GPU需要通訊的GPU個數就能降到9個，單次計算僅需900次。

實際上在此基礎上，Colossal-AI還包含2.5維、3維張量並行方法，可以進一步降低傳輸成本。

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

針對大圖片、視訊、長文字、長時間醫療監控等資料，Colossal-AI還使用了序列並行的方法，這種方法能突破原有機器能力限制，直接處理長序列資料。

值得一提的是，Colossal-AI的API介面是可以定製的，這使得它可以便捷新增新的並行維度。

其次，大規模優化器也是Colossal-AI的亮點。

上面我們也提到了，在分散式並行系統中會使用多種並行方法，資料並行則是另一種常見方法。

這種方法的原理不難理解，就是把訓練資料劃分成若干份，讓不同的機器運算不同的資料，然後通過一個引數伺服器 （Paremeter Server）收集目標資料。

由此可以大幅提升AI模型訓練過程中的批量大小，加速訓練過程。

不過大批量訓練有個“通病”，就是會產生泛化誤差 （Generalization Gap），導致網路泛化能力下降，進而導致AI模型準確度下降。

所以，Colossal-AI在系統中使用了自研的LAMB、LARS等大規模優化器。在保證訓練精度的情況下，還將批大小從512擴充套件到65536。

其中，LARS優化器是通過逐層調整學習率，來減少因為學習率導致的無法收斂情況。

LAMB優化器則是在LARS的基礎上，將逐層調整學習率的思想應用到自適應梯度上。

由此，LAMB能夠很好解決此前LARS在BERT訓練中存在差異的問題，最大批量達到了64K。

此前，LAMB優化器曾成功將預訓練一遍BERT的時間，從原本的三天三夜縮短到一個多小時。

第三方面，Colossal-AI使用自適應可擴充套件排程器來高效處理任務。

與現有常見的任務排程器不同，Colossal-AI不是靜態地通過GPU個數來判斷任務規模，而是根據批大小來動態、自動管理每個任務.

通過演化演算法，該任務排程器還能不斷優化排程決策，更大程度提升GPU利用率。

評估結果表明，與當前最先進的方法相比，該方法在平均JCT （job completion time）上能夠縮短45.6%的時間，優於現有的深度學習任務排程演算法。

此外，這種自適應可擴充套件排程器還能通過NCCL網路通訊實現高效的任務遷移。

最後，消除冗餘記憶體也是加速AI訓練的一種解決思路。

在這方面，Colossal-AI使用了zero redundancy optimizer技術（簡稱ZeRO）。

這種方法主要通過切分優化器狀態、梯度、模型引數，使GPU僅儲存當前計算所需的部分，從而來消除資料並行、模型並行中存在的記憶體冗餘。

尤其是在部署模型推理時，通過zero offload可以將模型解除安裝到CPU記憶體或硬碟，僅使用少量GPU資源，即可實現低成本部署前沿AI大模型。

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

而在應用層面，Colossal-AI的設計也顧及了能耗問題和易用性兩個維度。

考慮到資料移動會是能耗的主要來源，Colossal-AI在不增加計算量的情況下儘可能減少資料移動量，以此來降低能耗。

另一方面，作為一個開源給所有人使用的系統，Colossal-AI的使用門檻不高，即便是沒有學習過分散式系統的人也能上手操作。

同時，只需要極少量的程式碼改動，Colossal-AI就能將已有的單機程式碼快速擴充套件到平行計算叢集上。

最新實驗結果釋出

Talk is cheap，效果如何，還是得把實驗結果展開來看。

Colossal-AI近日釋出的最新實驗結果表明，這一大規模AI訓練系統具有通用性，在GPT-3、GPT-2、ViT、BERT等流行模型上均有亮眼的加速表現。

注：以下GPU均指英偉達A100。

GPT-3訓練速度提高10.7%

英偉達的Megatron-LM在加速訓練GPT-3時，至少需要128塊GPU才能啟動；而從下表可以看出，使用相同的計算資源，Colossal-AI可以將每次迭代花費的時間從43.1秒降至38.5秒。

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

站在工程的角度，考慮到訓練這樣的大模型往往需要投入數百萬美元，這一提升比例帶來的收益不言而喻。

另外，通過系統優化，Colossal-AI還能在訓練速度損失不大（43.1→48.5）的前提下，將GPU數量從128塊減少到96塊，大幅降低訓練成本。

而進一步啟用ZeRO3（零冗餘優化器）後，所需GPU數量甚至能減少一半——至64塊。

2天內可完成GPT-2訓練

在GPT-2的加速訓練結果中，可以看到，無論是在4、16還是64塊GPU的情況下，與Megatron-LM相比，Colossal-AI佔用的視訊記憶體都顯著減少。

也就是說，利用Colossal-AI，工程師們可以在採用同等數量GPU的前提下，訓練規模更大的模型，或設定更大的批量大小來加速訓練。

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

研發團隊在256塊GPU上進行了實驗，最終用時82.8個小時完成了15億引數版GPT-2的訓練。

據此預估，後續在512塊GPU上進行GPT-2預訓練，Colossal-AI能將訓練時間加速到45小時。

充分相容多種並行模式

在BERT上進行的實驗，則體現了Colossal-AI作為世界上並行維度最多的AI訓練系統的優勢。

與Megatron-LM相比，Colossal-AI序列並行方法只需要更少的視訊記憶體，就能夠利用更大的批量大小來加速訓練。同時，還允許開發者使用更長的序列資料。

Colossal-AI的序列並行方法還與流水並行方法相容。當開發者同時使用序列並行和流水並行時，可以進一步節省訓練大模型的時間。

另外，在近期的學術熱點ViT模型上，Colossal-AI也展現了高維張量並行方法的優勢。

在使用64張GPU的情況下，Colossal-AI採用2/2.5維方式進行張量並行，充分利用更大的批量大小，達到了更快的處理速度。

背後團隊：LAMB優化器作者尤洋領銜

看到這裡，是不是覺得Colossal-AI確實值得標星關注一發？

實際上，這一國產專案背後的研發團隊來頭不小。

領銜者，正是LAMB優化器的提出者尤洋。

在谷歌實習期間，正是憑藉LAMB，尤洋曾打破BERT預訓練世界紀錄。

據英偉達官方GitHub顯示，LAMB比Adam優化器快出整整72倍。微軟的DeepSpeed也採用了LAMB方法。

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

2020年博士畢業後，他加入新加坡國立大學計算機系，並於2021年1月成為校長青年教授（Presidential Young Professor）。

同樣是在2021年，他還獲得了IEEE-CS超算傑出新人獎。該獎項每年在全球範圍內表彰不超過3人，僅授予在博士畢業5年之內，已在高效能運算領域做出有影響力的卓越貢獻，並且可以為高效能運算的發展做出長期貢獻的優秀青年學者。

與此同時，尤洋回國創辦潞晨科技——一家主營業務為分散式軟體系統、大規模人工智慧平臺以及企業級雲端計算解決方案的AI初創公司。

其核心團隊成員來自加州大學伯克利分校、史丹佛大學、清華大學、北京大學、新加坡國立大學、新加坡南洋理工大學等國內外知名高校，在高效能運算、人工智慧、分散式系統方面有十餘年的技術積累，並已在國際頂級學術刊物/會議上發表論文30餘篇。

目前，潞晨科技已拿下創新工場和真格基金合投的超千萬元種子輪融資。

傳送門

有關Colossal-AI，今天就先介紹到這裡。

最後，附上傳送門，感興趣的小夥伴，自行取用~

GitHub地址：