在個人電腦用單塊 GPU 帶動180億引數 GPT！熱門開源專案再添新特性
提到訓練AI大模型，總能讓人想起動輒幾百上千塊GPU、天價訓練費用、只有幾家大廠才玩得起，普通AI玩家看著鋪天蓋地的大模型新聞只能默默流淚~
現在，僅有一塊GPU的個人PC也可以訓練高達180億引數GPT；普通的膝上型電腦，也能訓練十幾億引數的模型，相比現有主流方案，可提升引數容量十餘倍！如此顯著的提升來自Colossal-AI，一個通用AI大模型高效訓練系統。最重要的是，它完全開源，僅需極少量修改，即可讓現有深度學習專案在單張消費級顯示卡上使用大得多的模型進行訓練，每個人都可以在家訓練AI大模型！尤其是大幅度降低了AI大模型微調和推理等下游任務和應用部署的門檻！Colossal-AI還可將現有專案便捷擴充套件到大規模計算叢集，使用高效並行技術進一步加速。開源地址：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，美國科學院、工程院、藝術與科學院三院院士傳送門專案地址：https：//github。com/hpcaitech/ColossalAI參考連結：
https：//medium。com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc
