發表日期 4/6/2022, 5:00:35 PM
大規模並行AI訓練係統 Colossal-AI ,旨在作為深度學習框架的內核,幫助用戶便捷實現最大化提升AI部署效率,同時最小化部署成本。
開源地址:https://github.com/hpcaitech/ColossalAI
Colossal-AI一經開源便受到廣泛關注, 連續多日登頂GitHub熱榜Python方嚮世界第一 ,與眾多已有數萬star的明星開源項目一起受到海內外關注!
經過開發者們的不斷努力,Colossal-AI在數月的密集測試後迎來正式版!此版本由300多次commits組成。
本次正式版更新重點優化瞭 分布式訓練性能及開發者的易用性 ,主要亮點包括:
重構ZeRO以改善性能和易用性;添加細粒度Profiler TensorBoard監控插件,監測訓練過程中內存、網絡等狀態;更靈活的checkpoint策略,可擴展的pipeline模塊;開源蛋白質預測FastFold等豐富行業解決方案;添加中文教程,MOE、BERT等實例,開放用戶社群及論壇。
專業助力大模型訓練
近年來,隨著深度學習的興起及大模型橫掃各大性能榜單,前沿AI模型的大小在短短幾年內便已增大萬倍,遠超硬件數倍的緩慢增長。前沿AI大模型不僅遠超單個GPU的容納能力,所需算力也往往需要單個GPU運行數百甚至上韆年。
因此,如何提升單個GPU的容納能力,如何高效利用分布式技術,聯閤多個GPU低成本實現並行訓練加速已成為AI大模型的關鍵痛點。
針對現有方案並行維度有限、效率不高、通用性差、部署睏難、缺乏維護等痛點,Colossal-AI通過高效多維並行、顯存優化、大規模優化庫、細粒度監測等方式,讓用戶僅需極少量修改,即可高效快速部署AI大模型訓練。
多維並行
相比現有方案中並行維度僅包括數據並行、一維張量並行、流水並行三種方案,Colossal-AI進一步提供2/2.5/3維張量並行和序列並行,以及便捷的多維混閤並行解決方案。
△ ViT張量並行為64時,可提升14倍批大小和5倍訓練速度
其中,高維張量並行能極大減輕顯存消耗,提升通信效率,使得計算資源利用更加高效。
△ 序列並行幫助BERT提升2倍訓練速度,或1.5倍序列長度
而序列並行針對大圖片、視頻、長文本、長時間醫療監測等數據,可以幫助突破原有機器能力限製,直接處理長序列數據。
顯存優化
Colossal-AI綜閤瞭多重顯存優化技術,包含多維並行,ZeRO冗餘內存消除,CPU offload,Gradient Checkpoint,自動混閤精度(AMP)等前沿技術,最大限度幫助用戶避免顯存瓶頸,降低訓練的硬件需求。
△ GPT-2使用Colossal-AI,同樣硬件下提升24倍可訓練模型大小,或3倍訓練速度
靈活易用
Colossal-AI接口設計與PyTorch風格保持一緻,降低學習和使用成本,僅需極少量修改,便可將已有項目與Colossal-AI結閤,便捷擴展至大規模並行。此外,該係統還保持瞭優秀的擴展性,便於根據需求添加新功能,與已有功能模塊兼容。
細粒度監測 :細粒度Profiler TensorBoard插件,相較於PyTorch僅能以iteration為單位進行記錄訓練過程,Colossal-AI能夠監測iteration內的網絡、通信、內存等狀態,方便開發者進行精確分析和調試,提高開發效率。
大規模優化庫 :Colossal-AI提供大規模並行優化器LAMB、LARS等,首次將訓練batch size擴展到65536。Colossal-AI還與PyTorch自帶各類optimizer兼容,並不斷探索添加最新前沿優化技術,滿足各類模型需求。
豐富的行業解決方案
Colossal-AI目前已與 自動駕駛 、 雲計算 、 零售 、 醫藥 、 芯片 等行業知名廠商達成閤作,與AI領域頂級開源組織Hugging Face等建立閤作。
蛋白質結構預測加速方案:FastFold
AlphaFold因強大的AI預測蛋白質結構能力被Science和Nature評選為2021年十大科學突破之首,但存在訓練時間長、成本高等問題。
△ 圖片來源:https://arxiv.org/pdf/2203.00854.pdf
基於Colossal-AI的加速方案FastFold,將GPU優化和大模型訓練技術引入AlphaFold的訓練和推理,成功超越 榖歌 和 哥倫比亞大學 的方案,將AlphaFold訓練時間從11天減少到67小時,且總成本更低,在長序列推理中也實現9.3~11.6倍的速度提升。
△ 長序列推理性能對比
半數GPU訓練GPT-3
對於超大AI模型,如GPT-3,相比英偉達方案,Colossal-AI 僅需一半的計算資源 ,即可啓動訓練;若使用相同計算資源,則能提速11%,可 降低GPT-3訓練成本超百萬美元 。
Colossal-AI注重開源社區建設,提供中文教程,開放用戶社群及論壇,對於用戶反饋進行高效交流與迭代更新,不斷添加MoE等前沿應用。
項目團隊
潞晨技術團隊的核心成員均來自美國加州大學伯剋利分校,斯坦福大學,清華大學,北京大學,新加坡國立大學,新加坡南洋理工大學等國內外知名高校;擁有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名廠商工作經曆。公司成立即獲得創新工場、真格基金等多傢頂尖VC機構種子輪投資。
△ 潞晨科技創始人尤洋教授:加州大學伯剋利分校博士、IPDPS/ICPP最佳論文、ACM/IEEE George Michael HPC Fellowship、福布斯30歲以下精英(亞洲 2021)、IEEE-CS超算傑齣新人奬、UC伯剋利EECS Lotfi A. Zadeh優秀畢業生奬
△ 潞晨CSO Prof. James Demmel:加州大學伯剋利分校傑齣教授、ACM/IEEE Fellow,美國科學院、工程院、藝術與科學院三院院士
傳送門
論文地址:https://arxiv.org/abs/2110.14883
項目地址:https://github.com/hpcaitech/ColossalAI
文檔地址:https://www.colossalai.org/
*本文觀點參考鏈接:
https://medium.com/@hpcaitech/5-must-follow-features-that-are-seeing-colossal-ais-success-2d5361e27e4b
― 完 ―