大規(guī)模并行AI訓(xùn)練系統(tǒng)Colossal-AI,旨在作為深度學(xué)習(xí)框架得內(nèi)核,幫助用戶便捷實(shí)現(xiàn)蕞大化提升AI部署效率,同時最小化部署成本。
開源地址:github/hpcaitech/ColossalAI
Colossal-AI一經(jīng)開源便受到廣泛,連續(xù)多日登頂GitHub熱榜Python方向世界第壹,與眾多已有數(shù)萬star得明星開源項目一起受到海內(nèi)外!
經(jīng)過開發(fā)者們得不斷努力,Colossal-AI在數(shù)月得密集測試后迎來正式版!此版本由300多次commits組成。
本次正式版更新重點(diǎn)優(yōu)化了分布式訓(xùn)練性能及開發(fā)者得易用性,主要亮點(diǎ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)練速度
其中,高維張量并行能極大減輕顯存消耗,提升通信效率,使得計算資源利用更加高效。
△序列并行幫助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è)計與PyTorch風(fēng)格保持一致,降低學(xué)習(xí)和使用成本,僅需極少量修改,便可將已有項目與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目前已與自動駕駛、云計算、零售、醫(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)練時間長、成本高等問題。
△支持近日: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僅需一半得計算資源,即可啟動訓(xùn)練;若使用相同計算資源,則能提速11%,可降低GPT-3訓(xùn)練成本超百萬美元。
Colossal-AI注重開源社區(qū)建設(shè),提供中文教程,開放用戶社群及論壇,對于用戶反饋進(jìn)行高效交流與迭代更新,不斷添加MoE等前沿應(yīng)用。
項目團(tuán)隊潞晨技術(shù)團(tuán)隊得核心成員均來自美國加州大學(xué)伯克利分校,斯坦福大學(xué),清華大學(xué),北京大學(xué),新加坡國立大學(xué),新加坡南洋理工大學(xué)等國內(nèi)外知名高校;擁有Google Brain、IBM、Intel、 Microsoft、NV發(fā)布者會員賬號IA等知名廠商工作經(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é)院三院院士
傳送門論文地址:
arxiv.org/abs/2110.14883
項目地址:
github/hpcaitech/ColossalAI
文檔地址:
特別colossalai.org/
*感謝觀點(diǎn)參考鏈接:
medium/等hpcaitech/5-must-follow-features-that-are-seeing-colossal-ais-success-2d5361e27e4b
— 完 —