一個悄然崛起的國產(chǎn)開源項(xiàng)目
轉(zhuǎn)自 GitHubDaily
大家好，我是小 G。
此前，我曾在公眾號推薦過 GitHub 上一個性能頗為強(qiáng)悍的國產(chǎn) AI 開源項(xiàng)目：Colossal-AI，作為一個大規(guī)模并行 AI 訓(xùn)練系統(tǒng)、深度學(xué)習(xí)框架的內(nèi)核，該項(xiàng)目可幫助用戶便捷實(shí)現(xiàn)最大化提升 AI 部署效率，同時最小化部署成本。
僅需一半數(shù)量的 GPU，便能完成相同效果的 GPT-3 訓(xùn)練工作，極大降低了項(xiàng)目研發(fā)成本。
GitHub 地址：https://github.com/hpcaitech/ColossalAI
Colossal-AI 一經(jīng)開源便受到廣泛關(guān)注，連續(xù)多日登上 GitHub Python 方向世界第一的位置，與眾多已有數(shù)萬 star 的明星開源項(xiàng)目一起受到海內(nèi)外關(guān)注！
世界頂尖的人工智能機(jī)構(gòu) Hugging Face 也主動與 Colossal-AI 接洽合作，對于一個國產(chǎn)項(xiàng)目來說，受到如此關(guān)注，這也從側(cè)面反映出了它的實(shí)力。
近日，經(jīng)過開發(fā)者們的不斷努力后，Colossal-AI 在數(shù)月的密集測試后迎來正式版！此版本由 300 多次 commits 組成。
本次正式版更新重點(diǎn)優(yōu)化了分布式訓(xùn)練性能及開發(fā)者的易用性，主要亮點(diǎn)包括：
重構(gòu) ZeRO 以改善性能和易用性；
添加細(xì)粒度 Profiler TensorBoard 監(jiān)控插件，監(jiān)測訓(xùn)練過程中內(nèi)存、網(wǎng)絡(luò)等狀態(tài)；
更靈活的 checkpoint 策略，可擴(kuò)展的 pipeline 模塊；
開源蛋白質(zhì)預(yù)測 FastFold 等豐富行業(yè)解決方案；
添加中文教程，PaLM、MOE、BERT等實(shí)例，開放用戶社群及論壇。
專業(yè)助力大模型訓(xùn)練
近年來，隨著深度學(xué)習(xí)的興起及大模型橫掃各大性能榜單，前沿 AI 模型的大小在短短幾年內(nèi)便已增大萬倍，遠(yuǎn)超硬件數(shù)倍的緩慢增長。前沿 AI 大模型不僅遠(yuǎn)超單個 GPU 的容納能力，所需算力也往往需要單個 GPU 運(yùn)行數(shù)百甚至上千年。
因此，如何提升單個 GPU 的容納能力，如何高效利用分布式技術(shù)，聯(lián)合多個 GPU 低成本實(shí)現(xiàn)并行訓(xùn)練加速已成為 AI 大模型的關(guān)鍵痛點(diǎn)。
針對現(xiàn)有方案并行維度有限、效率不高、通用性差、部署困難、缺乏維護(hù)等痛點(diǎn)，Colossal-AI 通過高效多維并行、顯存優(yōu)化、大規(guī)模優(yōu)化庫、細(xì)粒度監(jiān)測等方式，讓用戶僅需極少量修改，即可高效快速部署 AI 大模型訓(xùn)練。
多維并行
相比現(xiàn)有方案中并行維度僅包括數(shù)據(jù)并行、一維張量并行、流水并行三種方案，Colossal-AI 進(jìn)一步提供 2/2.5/3 維張量并行和序列并行，以及便捷的多維混合并行解決方案。
△ViT 張量并行為 64 時，可提升 14 倍批大小和 5 倍訓(xùn)練速度
其中，高維張量并行能極大減輕顯存消耗，提升通信效率，使得計(jì)算資源利用更加高效。
△序列并行幫助 BERT 提升 2 倍訓(xùn)練速度，或 1.5 倍序列長度
而序列并行針對大圖片、視頻、長文本、長時間醫(yī)療監(jiān)測等數(shù)據(jù)，可以幫助突破原有機(jī)器能力限制，直接處理長序列數(shù)據(jù)。
顯存優(yōu)化
Colossal-AI 綜合了多重顯存優(yōu)化技術(shù)，包含多維并行，ZeRO 冗余內(nèi)存消除，CPU offload，Gradient Checkpoint，自動混合精度（AMP）等前沿技術(shù)，最大限度幫助用戶避免顯存瓶頸，降低訓(xùn)練的硬件需求。
△GPT-2 使用 Colossal-AI，同樣硬件下提升 24 倍可訓(xùn)練模型大小，或 3 倍訓(xùn)練速度
靈活易用
Colossal-AI 接口設(shè)計(jì)與 PyTorch 風(fēng)格保持一致，降低學(xué)習(xí)和使用成本，僅需極少量修改，便可將已有項(xiàng)目與 Colossal-AI 結(jié)合，便捷擴(kuò)展至大規(guī)模并行。此外，該系統(tǒng)還保持了優(yōu)秀的擴(kuò)展性，便于根據(jù)需求添加新功能，與已有功能模塊兼容。
細(xì)粒度監(jiān)測：細(xì)粒度 Profiler TensorBoard 插件，相較于 PyTorch 僅能以 iteration 為單位進(jìn)行記錄訓(xùn)練過程，Colossal-AI 能夠監(jiān)測 iteration 內(nèi)的網(wǎng)絡(luò)、通信、內(nèi)存等狀態(tài)，方便開發(fā)者進(jìn)行精確分析和調(diào)試，提高開發(fā)效率。
大規(guī)模優(yōu)化庫：Colossal-AI 提供大規(guī)模并行優(yōu)化器 LAMB、LARS 等，首次將訓(xùn)練 batch size 擴(kuò)展到 65536。Colossal-AI 還與 PyTorch 自帶各類 optimizer 兼容，并不斷探索添加最新前沿優(yōu)化技術(shù)，滿足各類模型需求。
豐富的行業(yè)解決方案
Colossal-AI 目前已與自動駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)知名廠商達(dá)成合作，與 AI 領(lǐng)域頂級開源組織 Hugging Face 等建立合作。
蛋白質(zhì)結(jié)構(gòu)預(yù)測加速方案：FastFold
AlphaFold 因強(qiáng)大的 AI 預(yù)測蛋白質(zhì)結(jié)構(gòu)能力被 Science 和 Nature 評選為 2021 年十大科學(xué)突破之首，但存在訓(xùn)練時間長、成本高等問題。
△圖片來源：https://arxiv.org/pdf/2203.00854.pdf
基于 Colossal-AI 的加速方案 FastFold，將 GPU 優(yōu)化和大模型訓(xùn)練技術(shù)引入 AlphaFold 的訓(xùn)練和推理，成功超越谷歌和哥倫比亞大學(xué)的方案，將 AlphaFold 訓(xùn)練時間從 11 天減少到 67 小時，且總成本更低，在長序列推理中也實(shí)現(xiàn) 9.3～11.6 倍的速度提升。
△長序列推理性能對比
半數(shù) GPU 訓(xùn)練 GPT-3
對于超大 AI 模型，如 GPT-3，相比英偉達(dá)方案，Colossal-AI 僅需一半的計(jì)算資源，即可啟動訓(xùn)練；若使用相同計(jì)算資源，則能提速 11%，可降低 GPT-3 訓(xùn)練成本超百萬美元。
Colossal-AI 注重開源社區(qū)建設(shè)，提供中文教程，開放用戶社群及論壇，對于用戶反饋進(jìn)行高效交流與迭代更新，不斷添加 PaLM、MOE 等前沿應(yīng)用。
項(xiàng)目團(tuán)隊(duì)
潞晨技術(shù)團(tuán)隊(duì)的核心成員均來自美國加州大學(xué)伯克利分校，斯坦福大學(xué)，清華大學(xué)，北京大學(xué)，新加坡國立大學(xué)，新加坡南洋理工大學(xué)等國內(nèi)外知名高校；擁有 Google Brain、IBM、Intel、 Microsoft、NVIDIA 等知名廠商工作經(jīng)歷。公司成立即獲得創(chuàng)新工場、真格基金等多家頂尖 VC 機(jī)構(gòu)種子輪投資。
△潞晨科技創(chuàng)始人尤洋教授：加州大學(xué)伯克利分校博士、IPDPS/ICPP 最佳論文、ACM/IEEE George Michael HPC Fellowship、福布斯 30 歲以下精英 (亞洲 2021)、IEEE-CS 超算杰出新人獎、UC 伯克利 EECS Lotfi A. Zadeh 優(yōu)秀畢業(yè)生獎
△潞晨 CSO Prof. James Demmel：加州大學(xué)伯克利分校杰出教授、ACM/IEEE Fellow，美國科學(xué)院、工程院、藝術(shù)與科學(xué)院三院院士
傳送門
論文地址：https://arxiv.org/abs/2110.14883
項(xiàng)目地址：https://github.com/hpcaitech/ColossalAI
文檔地址：https://www.colossalai.org/
* 本文觀點(diǎn)參考鏈接：
https://medium.com/@hpcaitech/5-must-follow-features-that-are-seeing-colossal-ais-success-2d5361e27e4b