一個悄然崛起的國產開源項目 Colossal-AI
大家好，我是小 G。
提到訓練 AI 大模型，總能讓人想起動輒幾百上千塊 GPU、天價訓練費用、只有幾家大廠才玩得起，普通 AI 玩家看着鋪天蓋地的大模型新聞只能默默流淚～
現在，僅有一塊 GPU 的個人 PC 也可以訓練高達 180 億參數 GPT；普通的筆記本電腦，也能訓練十幾億參數的模型，相比現有主流方案，可提升參數容量十餘倍！
如此顯著的提升來自 GitHub 上的一個國產項目：Colossal-AI，一個通用 AI 大模型高效訓練系統。最重要的是，它完全開源，僅需極少量修改，即可讓現有深度學習項目在單張消費級顯卡上使用大得多的模型進行訓練，每個人都可以在家訓練 AI 大模型！尤其是大幅度降低了 AI 大模型微調和推理等下游任務和應用部署的門檻！
Colossal-AI 還可將現有項目便捷擴展到大規模計算集羣，使用高效並行技術進一步加速。
GitHub 地址：https://github.com/hpcaitech/ColossalAI
巨頭角力，爭煉 AI 大模型
從 2018 年穀歌提出的 3 億參數 BERT 起，大模型記錄在短短几年時間內被不斷刷新，OpenAI 1750 億參數的 GPT-3，微軟和英偉達聯手發佈的 5300 億參數 MT-NLG ......
稠密單體模型已達千億參數，而稀疏混合模型，如谷歌 2021 年發佈的 Switch Transformer，更是將參數量推至萬億級別。
然而，從頭訓練如此之大的模型費用極高，通常需要同時使用數百甚至上千張 NVIDIA A100 等專業高性能計算 GPU，使用專用的 InfiniBand 高速網絡互聯組成超級計算機集羣，單次訓練成本可達千萬美元。
使用單張消費級顯卡訓練 AI 大模型
對於在校學生和個人開發者等普通 AI 玩家來說，顯然無法承擔上述如此高昂的費用，最方便獲得的計算資源還是個人電腦中的英偉達 RTX 系列消費級 GPU。
爲了解放 AI 生產力，讓大模型普惠到更多開發人員，真正實現 "多快好省" 的使用 AI 大模型，Colossal-AI 僅需幾行代碼，便可實現提升單卡訓練模型容量十幾倍的提升。
在各類型硬件上，Colossal-AI 的表現均遠超原生 PyTorch 和主流的微軟 DeepSpeed 方案。
對於大模型的代表 GPT，使用 Colossal-AI 在搭載 RTX 2060 6GB 的普通遊戲筆記本上，也足以訓練高達 15 億參數模型；對於搭載 RTX3090 24GB 的個人電腦，更是可以直接訓練 180 億參數的模型；對於 Tesla V100 等專業計算卡，Colossal-AI 也能顯示出顯著改善。
Colossal-AI 還成功復現了谷歌的最新研究成果 PaLM (Pathways Language Model)，在各類硬件上也體現出卓越的性能提升，而微軟 DeepSpeed 尚不支持 PaLM 模型。
關鍵技術：異構訓練再升級
使用單張消費級顯卡訓練 AI 大模型的最大困難在於顯存容量極其有限，嚴重限制了可容納的模型參數量。微軟 DeepSpeed 的提出 ZeRO-offload 方法，嘗試將模型切分，利用更大容量、更低成本的內存。目前已經出現了多個基於 DeepSpeed 的魔改版本。但如下圖左邊所示，當 GPU 內存不足以滿足其相應的模型數據要求時，即使當時 CPU 上仍有可用內存，系統也會崩潰。
不同於基於 DeepSpeed 的各種魔改方案，Colossal-AI 團隊從頭搭建瞭如 ZeRO 等核心關鍵技術，並針對 DeepSpeed 在 CPU 和 GPU 內存之間僅使用靜態劃分模型數據、對不同訓練配置使用固定內存佈局等問題做了諸多改進，進一步挖掘高效的 GPU 與 CPU 內存高效協同方案，畢竟內存條可比大顯存的高端顯卡便宜太多～
Colossal-AI 設計的 Gemini，就像雙子星一樣，高效管理和利用 GPU 與 CPU 的異構內存，讓張量在訓練過程中動態分佈在 CPU-GPU 的存儲空間內，從而讓模型訓練突破 GPU 的內存牆。
我們利用深度學習網絡訓練過程的迭代特性，按迭代次數將訓練分爲 warmup 和 non-warmup 兩個階段。在初期 warmup 階段，監測內存信息；在 non-warmup 階段利用已收集的信息來高效移動張量，以達到最小化 CPU-GPU 數據移動的目的。
聽起來容易，實現起來卻並不簡單。非模型的內存使用量其實難以獲取，因爲非模型數據的生存週期並不歸用戶管理，現有的深度學習框架沒有暴露非模型數據的追蹤接口給用戶。其次，CUDA context 等非框架開銷也需要統計。
Colossal-AI 通過採樣方式在 warmup 階段獲得 CPU 和 GPU 內存的使用情況。非模型數據的使用可以通過統計兩個時刻之間系統最大內存使用 - 模型內存使用獲得。模型的內存使用情況可以通過查詢內存管理器得知，如下圖黑色實線所示。
而所有模型數據張量則交給內存管理器管理，每個張量標記一個狀態信息，包括 HOLD，COMPUTE，FREE 等。並根據動態查詢到的內存使用情況，不斷動態轉換張量狀態，調整張量位置，最終實現對 GPU 顯存和 CPU 內存的高效利用，實現在硬件極其有限的情況下，最大化模型容量和平衡訓練速度，對於 AI 民主化和低成本微調大模型下游任務等意義巨大。
更進一步：便捷高效並行擴展
並行分佈式技術是進一步加速模型訓練的重要手段，想要以最短時間訓練當今世界最大最前沿的 AI 模型，仍離不開高效的分佈式並行擴展。針對現有方案並行維度有限、效率不高、通用性差、部署困難、缺乏維護等痛點，Colossal-AI 通過高效多維並行和異構並行等技術，讓用戶僅需極少量修改，即可高效快速部署 AI 大模型訓練。
例如對於 GPT-3 這樣的超大 AI 模型，相比英偉達方案，Colossal-AI 僅需一半的計算資源，即可啓動訓練；若使用相同計算資源，則能提速 11%，可降低 GPT-3 訓練成本超百萬美元。
對於蛋白質結構預測應用 AlphaFold，基於 Colossal-AI 的加速方案的 FastFold，成功超越谷歌和哥倫比亞大學的方案，將 AlphaFold 訓練時間從 11 天減少到 67 小時，且總成本更低，在長序列推理中也實現 9.3～11.6 倍的速度提升。
Colossal-AI 注重開源社區建設，提供中文教程，開放用戶社羣及論壇，對於用戶反饋進行高效交流與迭代更新，不斷添加 PaLM、AlphaFold 等前沿應用。
自開源以來，Colossal-AI 已經多次登頂 GitHub 熱榜 Python 方向世界第一，與衆多已有數萬 star 的明星開源項目一起受到海內外關注！
項目團隊
潞晨技術團隊的核心成員均來自美國加州大學伯克利分校，斯坦福大學，清華大學，北京大學，新加坡國立大學，新加坡南洋理工大學等國內外知名高校；擁有 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，美國科學院、工程院、藝術與科學院三院院士
傳送門
GitHub 地址：https://github.com/hpcaitech/ColossalAI
參考鏈接：
https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc
GitHub 寶典 專注於分享 GitHub 上知名的 Python、Java、Web、AI、數據分析等多個領域的優質學習資源、開源項目及開發者工具，爲 GitHub 開發者提供更好的資訊服務。
本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/exEMTa-UN8ZwQgYvyOH6JQ