1 月 20 日消息，當(dāng)今 AI 之勢，影響縱深發(fā)展的矛盾是什么？
一方面，大模型風(fēng)頭正勁，效果驚艷，人人都想試試。但另一方面，硬件基礎(chǔ)上動不動就是上萬張 GPU 的大規(guī)模集群在日夜燃燒，鈔能力勸退。
所以如果告訴你，現(xiàn)在只用一半數(shù)量的 GPU，也能完成同樣的 GPT-3 訓(xùn)練呢？
你會覺得關(guān)鍵鑰匙是什么？
不賣關(guān)子了。實現(xiàn)如此提升的，是一個名為 Colossal-AI 的 GitHub 開源項目。
而且該項目開源不久，就迅速登上了 Python 方向的熱榜世界第一。
GitHub 地址：點擊打開
不僅能加速 GPT-3，對于 GPT-2、ViT、BERT 等多種模型，Colossal-AI 的表現(xiàn)也都非常 nice：
比如半小時左右就能預(yù)訓(xùn)練一遍 ViT-Base / 32，2 天能訓(xùn)完 15 億參數(shù) GPT 模型、5 天可訓(xùn)完 83 億參數(shù) GPT 模型。
與業(yè)內(nèi)主流的 AI 并行系統(tǒng) —— 英偉達 Megatron-LM 相比，在同樣使用 512 塊 GPU 訓(xùn)練 GPT-2 模型時，Colossal-AI 的加速比是其 2 倍。而在訓(xùn)練 GPT-3 時，更是可以節(jié)省近千萬元的訓(xùn)練費用。
此外在訓(xùn)練 GPT-2 時，顯存消耗甚至能控制在 Megatron-LM 的十分之一以下。
Colossal-AI 究竟是如何做到的？
老規(guī)矩，我們從論文扒起。
高效 6 維并行方法
簡單來說，Colossal-AI 就是一個整合了多種并行方法的系統(tǒng)，提供的功能包括多維并行、大規(guī)模優(yōu)化器、自適應(yīng)任務(wù)調(diào)度、消除冗余內(nèi)存等。
首先來看多維并行。
所謂“多維”是指，目前主流的分布式并行方案往往使用多種并行方法。
比如英偉達的 Megatron-LM 使用了 3 種方法：數(shù)據(jù)并行、流水并行和張量并行。因此這種模式也被稱為三維并行。微軟的 DeepSpeed 調(diào)用 Megatron-LM 作為并行基礎(chǔ)。
而 Colossal-AI 能將系統(tǒng)的并行維度，一下子拉升到 6 維 ——
在兼容數(shù)據(jù)并行、流水并行的基礎(chǔ)上，基于該項目團隊自研的 2 維 / 2.5 維 / 3 維張量并行方法，以及序列并行實現(xiàn)。
其中，高維張量并行正是 Colossal-AI 提升大模型顯存利用率和通信效率的關(guān)鍵所在。
其實張量并行并不新奇，只是過去我們常見的張量并行更多都是基于一維的。
它的原理是將模型層內(nèi)的權(quán)重參數(shù)按行或列切分到不同的處理器上，利用分塊矩陣乘法，將一個運算分布到多個處理器上同時進行。
比如英偉達的 Megatron-LM 就是一個典型的例子。
但這種并行方式存在一定弊端。
比如，每個處理器仍需要存儲整個中間激活，使得在處理大模型時會浪費大量顯存空間。
另一方面，這種單線方法還會導(dǎo)致每個處理器都需要與其他所有處理器進行通信。
這意味著假設(shè)有 100 個 GPU 的話，每個 GPU 都需要與其他 99 個 GPU 通信，每次計算需要通信的次數(shù)就高達 9900 次。
但如果將張量并行的維度擴展到 2 維，單次計算量能立刻下降一個量級。
因為每個 GPU 只需與自己同行或同列的 GPU 通信即可。
同樣還是 100 個 GPU 的情況，每個 GPU 需要通信的 GPU 個數(shù)就能降到 9 個，單次計算僅需 900 次。
實際上在此基礎(chǔ)上，Colossal-AI 還包含 2.5 維、3 維張量并行方法，可以進一步降低傳輸成本。
相較于 2 維并行方法，2.5 維并行方法可提升 1.45 倍效率，3 維方法可提升 1.57 倍。
針對大圖片、視頻、長文本、長時間醫(yī)療監(jiān)控等數(shù)據(jù)，Colossal-AI 還使用了序列并行的方法，這種方法能突破原有機器能力限制，直接處理長序列數(shù)據(jù)。
值得一提的是，Colossal-AI 的 API 接口是可以定制的，這使得它可以便捷添加新的并行維度。
其次，大規(guī)模優(yōu)化器也是 Colossal-AI 的亮點。
上面我們也提到了，在分布式并行系統(tǒng)中會使用多種并行方法，數(shù)據(jù)并行則是另一種常見方法。
這種方法的原理不難理解，就是把訓(xùn)練數(shù)據(jù)劃分成若干份，讓不同的機器運算不同的數(shù)據(jù)，然后通過一個參數(shù)服務(wù)器 （Paremeter Server）收集目標(biāo)數(shù)據(jù)。
由此可以大幅提升 AI 模型訓(xùn)練過程中的批量大小，加速訓(xùn)練過程。
不過大批量訓(xùn)練有個“通病”，就是會產(chǎn)生泛化誤差 （Generalization Gap），導(dǎo)致網(wǎng)絡(luò)泛化能力下降，進而導(dǎo)致 AI 模型準(zhǔn)確度下降。
所以，Colossal-AI 在系統(tǒng)中使用了自研的 LAMB、LARS 等大規(guī)模優(yōu)化器。在保證訓(xùn)練精度的情況下，還將批大小從 512 擴展到 65536。
其中，LARS 優(yōu)化器是通過逐層調(diào)整學(xué)習(xí)率，來減少因為學(xué)習(xí)率導(dǎo)致的無法收斂情況。
LAMB 優(yōu)化器則是在 LARS 的基礎(chǔ)上，將逐層調(diào)整學(xué)習(xí)率的思想應(yīng)用到自適應(yīng)梯度上。
由此，LAMB 能夠很好解決此前 LARS 在 BERT 訓(xùn)練中存在差異的問題，最大批量達到了 64K。
此前，LAMB 優(yōu)化器曾成功將預(yù)訓(xùn)練一遍 BERT 的時間，從原本的三天三夜縮短到一個多小時。
第三方面，Colossal-AI 使用自適應(yīng)可擴展調(diào)度器來高效處理任務(wù)。
與現(xiàn)有常見的任務(wù)調(diào)度器不同，Colossal-AI 不是靜態(tài)地通過 GPU 個數(shù)來判斷任務(wù)規(guī)模，而是根據(jù)批大小來動態(tài)、自動管理每個任務(wù).
通過演化算法，該任務(wù)調(diào)度器還能不斷優(yōu)化調(diào)度決策，更大程度提升 GPU 利用率。
評估結(jié)果表明，與當(dāng)前最先進的方法相比，該方法在平均 JCT （job completion time）上能夠縮短 45.6% 的時間，優(yōu)于現(xiàn)有的深度學(xué)習(xí)任務(wù)調(diào)度算法。
此外，這種自適應(yīng)可擴展調(diào)度器還能通過 NCCL 網(wǎng)絡(luò)通信實現(xiàn)高效的任務(wù)遷移。
最后，消除冗余內(nèi)存也是加速 AI 訓(xùn)練的一種解決思路。
在這方面，Colossal-AI 使用了 zero redundancy optimizer 技術(shù)（簡稱 ZeRO）。
這種方法主要通過切分優(yōu)化器狀態(tài)、梯度、模型參數(shù)，使 GPU 僅保存當(dāng)前計算所需的部分，從而來消除數(shù)據(jù)并行、模型并行中存在的內(nèi)存冗余。
尤其是在部署模型推理時，通過 zero offload 可以將模型卸載到 CPU 內(nèi)存或硬盤，僅使用少量 GPU 資源，即可實現(xiàn)低成本部署前沿 AI 大模型。
綜上不難看出，在技術(shù)層面 Colossal-AI 的加速效果非常明顯。
而在應(yīng)用層面，Colossal-AI 的設(shè)計也顧及了能耗問題和易用性兩個維度。
考慮到數(shù)據(jù)移動會是能耗的主要來源，Colossal-AI 在不增加計算量的情況下盡可能減少數(shù)據(jù)移動量，以此來降低能耗。
另一方面，作為一個開源給所有人使用的系統(tǒng)，Colossal-AI 的使用門檻不高，即便是沒有學(xué)習(xí)過分布式系統(tǒng)的人也能上手操作。
同時，只需要極少量的代碼改動，Colossal-AI 就能將已有的單機代碼快速擴展到并行計算集群上。
最新實驗結(jié)果釋出
Talk is cheap，效果如何，還是得把實驗結(jié)果展開來看。
Colossal-AI 近日釋出的最新實驗結(jié)果表明，這一大規(guī)模 AI 訓(xùn)練系統(tǒng)具有通用性，在 GPT-3、GPT-2、ViT、BERT 等流行模型上均有亮眼的加速表現(xiàn)。
注：以下 GPU 均指英偉達 A100。
GPT-3 訓(xùn)練速度提高 10.7%
英偉達的 Megatron-LM 在加速訓(xùn)練 GPT-3 時，至少需要 128 塊 GPU 才能啟動；而從下表可以看出，使用相同的計算資源，Colossal-AI 可以將每次迭代花費的時間從 43.1 秒降至 38.5 秒。
這也就意味著，Colossal-AI 可以將 GPT-3 的訓(xùn)練速度進一步提高 10.7%。
站在工程的角度，考慮到訓(xùn)練這樣的大模型往往需要投入數(shù)百萬美元，這一提升比例帶來的收益不言而喻。
另外，通過系統(tǒng)優(yōu)化，Colossal-AI 還能在訓(xùn)練速度損失不大（43.1→48.5）的前提下，將 GPU 數(shù)量從 128 塊減少到 96 塊，大幅降低訓(xùn)練成本。
而進一步啟用 ZeRO3（零冗余優(yōu)化器）后，所需 GPU 數(shù)量甚至能減少一半 —— 至 64 塊。
2 天內(nèi)可完成 GPT-2 訓(xùn)練
在 GPT-2 的加速訓(xùn)練結(jié)果中，可以看到，無論是在 4、16 還是 64 塊 GPU 的情況下，與 Megatron-LM 相比，Colossal-AI 占用的顯存都顯著減少。
也就是說，利用 Colossal-AI，工程師們可以在采用同等數(shù)量 GPU 的前提下，訓(xùn)練規(guī)模更大的模型，或設(shè)置更大的批量大小來加速訓(xùn)練。
從下表結(jié)果中還可以看出，隨著批量大小的增加，Colossal-AI 的資源利用率會進一步提高，達到 Megatron-LM 速度的 2 倍。
研發(fā)團隊在 256 塊 GPU 上進行了實驗，最終用時 82.8 個小時完成了 15 億參數(shù)版 GPT-2 的訓(xùn)練。
據(jù)此預(yù)估，后續(xù)在 512 塊 GPU 上進行 GPT-2 預(yù)訓(xùn)練，Colossal-AI 能將訓(xùn)練時間加速到 45 小時。
充分兼容多種并行模式
在 BERT 上進行的實驗，則體現(xiàn)了 Colossal-AI 作為世界上并行維度最多的 AI 訓(xùn)練系統(tǒng)的優(yōu)勢。
與 Megatron-LM 相比，Colossal-AI 序列并行方法只需要更少的顯存，就能夠利用更大的批量大小來加速訓(xùn)練。同時，還允許開發(fā)者使用更長的序列數(shù)據(jù)。
Colossal-AI 的序列并行方法還與流水并行方法兼容。當(dāng)開發(fā)者同時使用序列并行和流水并行時，可以進一步節(jié)省訓(xùn)練大模型的時間。
另外，在近期的學(xué)術(shù)熱點 ViT 模型上，Colossal-AI 也展現(xiàn)了高維張量并行方法的優(yōu)勢。
在使用 64 張 GPU 的情況下，Colossal-AI 采用 2/2.5 維方式進行張量并行，充分利用更大的批量大小，達到了更快的處理速度。
背后團隊：LAMB 優(yōu)化器作者尤洋領(lǐng)銜
看到這里，是不是覺得 Colossal-AI 確實值得標(biāo)星關(guān)注一發(fā)？
實際上，這一國產(chǎn)項目背后的研發(fā)團隊來頭不小。
領(lǐng)銜者，正是 LAMB 優(yōu)化器的提出者尤洋。
在谷歌實習(xí)期間，正是憑借 LAMB，尤洋曾打破 BERT 預(yù)訓(xùn)練世界紀(jì)錄。
據(jù)英偉達官方 GitHub 顯示，LAMB 比 Adam 優(yōu)化器快出整整 72 倍。微軟的 DeepSpeed 也采用了 LAMB 方法。
說回到尤洋本人，他曾以第一名的成績保送清華計算機系碩士研究生，后赴加州大學(xué)伯克利分校攻讀 CS 博士學(xué)位。
2020 年博士畢業(yè)后，他加入新加坡國立大學(xué)計算機系，并于 2021 年 1 月成為校長青年教授（Presidential Young Professor）。
同樣是在 2021 年，他還獲得了 IEEE-CS 超算杰出新人獎。該獎項每年在全球范圍內(nèi)表彰不超過 3 人，僅授予在博士畢業(yè) 5 年之內(nèi)，已在高性能計算領(lǐng)域做出有影響力的卓越貢獻，并且可以為高性能計算的發(fā)展做出長期貢獻的優(yōu)秀青年學(xué)者。
與此同時，尤洋回國創(chuàng)辦潞晨科技 —— 一家主營業(yè)務(wù)為分布式軟件系統(tǒng)、大規(guī)模人工智能平臺以及企業(yè)級云計算解決方案的 AI 初創(chuàng)公司。
其核心團隊成員來自加州大學(xué)伯克利分校、斯坦福大學(xué)、清華大學(xué)、北京大學(xué)、新加坡國立大學(xué)、新加坡南洋理工大學(xué)等國內(nèi)外知名高校，在高性能計算、人工智能、分布式系統(tǒng)方面有十余年的技術(shù)積累，并已在國際頂級學(xué)術(shù)刊物 / 會議上發(fā)表論文 30 余篇。
目前，潞晨科技已拿下創(chuàng)新工場和真格基金合投的超千萬元種子輪融資。
傳送門
有關(guān) Colossal-AI，今天就先介紹到這里。
最后，附上傳送門，感興趣的小伙伴，自行取用~
GitHub 地址：點擊打開
參考鏈接：點擊打開
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。