提到訓(xùn)練 AI 大模型，總能讓人想起動(dòng)輒幾百上千塊 GPU、天價(jià)訓(xùn)練費(fèi)用、只有幾家大廠才玩得起，普通 AI 玩家看著鋪天蓋地的大模型新聞只能默默流淚～現(xiàn)在，僅有一塊 GPU 的個(gè)人 PC 也可以訓(xùn)練高達(dá) 180 億參數(shù) GPT；普通的筆記本電腦，也能訓(xùn)練十幾億參數(shù)的模型，相比現(xiàn)有主流方案，可提升參數(shù)容量十余倍！如此顯著的提升來自 GitHub 上的一個(gè)國產(chǎn)項(xiàng)目：Colossal-AI，一個(gè)通用 AI 大模型高效訓(xùn)練系統(tǒng)。最重要的是，它完全開源，僅需極少量修改，即可讓現(xiàn)有深度學(xué)習(xí)項(xiàng)目在單張消費(fèi)級顯卡上使用大得多的模型進(jìn)行訓(xùn)練，每個(gè)人都可以在家訓(xùn)練 AI 大模型！尤其是大幅度降低了 AI 大模型微調(diào)和推理等下游任務(wù)和應(yīng)用部署的門檻！Colossal-AI 還可將現(xiàn)有項(xiàng)目便捷擴(kuò)展到大規(guī)模計(jì)算集群，使用高效并行技術(shù)進(jìn)一步加速。GitHub 地址：https://github.com/hpcaitech/ColossalAI從 2018 年谷歌提出的 3 億參數(shù) BERT 起，大模型記錄在短短幾年時(shí)間內(nèi)被不斷刷新，OpenAI 1750 億參數(shù)的 GPT-3，微軟和英偉達(dá)聯(lián)手發(fā)布的 5300 億參數(shù) MT-NLG ......
稠密單體模型已達(dá)千億參數(shù)，而稀疏混合模型，如谷歌 2021 年發(fā)布的 Switch Transformer，更是將參數(shù)量推至萬億級別。然而，從頭訓(xùn)練如此之大的模型費(fèi)用極高，通常需要同時(shí)使用數(shù)百甚至上千張 NVIDIA A100 等專業(yè)高性能計(jì)算 GPU，使用專用的 InfiniBand 高速網(wǎng)絡(luò)互聯(lián)組成超級計(jì)算機(jī)集群，單次訓(xùn)練成本可達(dá)千萬美元。使用單張消費(fèi)級顯卡訓(xùn)練 AI 大模型對于在校學(xué)生和個(gè)人開發(fā)者等普通 AI 玩家來說，顯然無法承擔(dān)上述如此高昂的費(fèi)用，最方便獲得的計(jì)算資源還是個(gè)人電腦中的英偉達(dá) RTX 系列消費(fèi)級 GPU。為了解放 AI 生產(chǎn)力，讓大模型普惠到更多開發(fā)人員，真正實(shí)現(xiàn) '多快好省' 的使用 AI 大模型，Colossal-AI 僅需幾行代碼，便可實(shí)現(xiàn)提升單卡訓(xùn)練模型容量十幾倍的提升。在各類型硬件上，Colossal-AI 的表現(xiàn)均遠(yuǎn)超原生 PyTorch 和主流的微軟 DeepSpeed 方案。對于大模型的代表 GPT，使用 Colossal-AI 在搭載 RTX 2060 6GB 的普通游戲筆記本上，也足以訓(xùn)練高達(dá) 15 億參數(shù)模型；對于搭載 RTX3090 24GB 的個(gè)人電腦，更是可以直接訓(xùn)練 180 億參數(shù)的模型；對于 Tesla V100 等專業(yè)計(jì)算卡，Colossal-AI 也能顯示出顯著改善。Colossal-AI 還成功復(fù)現(xiàn)了谷歌的最新研究成果 PaLM (Pathways Language Model)，在各類硬件上也體現(xiàn)出卓越的性能提升，而微軟 DeepSpeed 尚不支持 PaLM 模型。關(guān)鍵技術(shù)：異構(gòu)訓(xùn)練再升級使用單張消費(fèi)級顯卡訓(xùn)練 AI 大模型的最大困難在于顯存容量極其有限，嚴(yán)重限制了可容納的模型參數(shù)量。微軟 DeepSpeed 的提出 ZeRO-offload 方法，嘗試將模型切分，利用更大容量、更低成本的內(nèi)存。目前已經(jīng)出現(xiàn)了多個(gè)基于 DeepSpeed 的魔改版本。但如下圖左邊所示，當(dāng) GPU 內(nèi)存不足以滿足其相應(yīng)的模型數(shù)據(jù)要求時(shí)，即使當(dāng)時(shí) CPU 上仍有可用內(nèi)存，系統(tǒng)也會(huì)崩潰。不同于基于 DeepSpeed 的各種魔改方案，Colossal-AI 團(tuán)隊(duì)從頭搭建了如 ZeRO 等核心關(guān)鍵技術(shù)，并針對 DeepSpeed 在 CPU 和 GPU 內(nèi)存之間僅使用靜態(tài)劃分模型數(shù)據(jù)、對不同訓(xùn)練配置使用固定內(nèi)存布局等問題做了諸多改進(jìn)，進(jìn)一步挖掘高效的 GPU 與 CPU 內(nèi)存高效協(xié)同方案，畢竟內(nèi)存條可比大顯存的高端顯卡便宜太多～Colossal-AI 設(shè)計(jì)的 Gemini，就像雙子星一樣，高效管理和利用 GPU 與 CPU 的異構(gòu)內(nèi)存，讓張量在訓(xùn)練過程中動(dòng)態(tài)分布在 CPU-GPU 的存儲(chǔ)空間內(nèi)，從而讓模型訓(xùn)練突破 GPU 的內(nèi)存墻。我們利用深度學(xué)習(xí)網(wǎng)絡(luò)訓(xùn)練過程的迭代特性，按迭代次數(shù)將訓(xùn)練分為 warmup 和 non-warmup 兩個(gè)階段。在初期 warmup 階段，監(jiān)測內(nèi)存信息；在 non-warmup 階段利用已收集的信息來高效移動(dòng)張量，以達(dá)到最小化 CPU-GPU 數(shù)據(jù)移動(dòng)的目的。聽起來容易，實(shí)現(xiàn)起來卻并不簡單。非模型的內(nèi)存使用量其實(shí)難以獲取，因?yàn)榉悄Ｐ蛿?shù)據(jù)的生存周期并不歸用戶管理，現(xiàn)有的深度學(xué)習(xí)框架沒有暴露非模型數(shù)據(jù)的追蹤接口給用戶。其次，CUDA context 等非框架開銷也需要統(tǒng)計(jì)。Colossal-AI 通過采樣方式在 warmup 階段獲得 CPU 和 GPU 內(nèi)存的使用情況。非模型數(shù)據(jù)的使用可以通過統(tǒng)計(jì)兩個(gè)時(shí)刻之間系統(tǒng)最大內(nèi)存使用 - 模型內(nèi)存使用獲得。模型的內(nèi)存使用情況可以通過查詢內(nèi)存管理器得知，如下圖黑色實(shí)線所示。而所有模型數(shù)據(jù)張量則交給內(nèi)存管理器管理，每個(gè)張量標(biāo)記一個(gè)狀態(tài)信息，包括 HOLD，COMPUTE，F(xiàn)REE 等。并根據(jù)動(dòng)態(tài)查詢到的內(nèi)存使用情況，不斷動(dòng)態(tài)轉(zhuǎn)換張量狀態(tài)，調(diào)整張量位置，最終實(shí)現(xiàn)對 GPU 顯存和 CPU 內(nèi)存的高效利用，實(shí)現(xiàn)在硬件極其有限的情況下，最大化模型容量和平衡訓(xùn)練速度，對于 AI 民主化和低成本微調(diào)大模型下游任務(wù)等意義巨大。更進(jìn)一步：便捷高效并行擴(kuò)展
并行分布式技術(shù)是進(jìn)一步加速模型訓(xùn)練的重要手段，想要以最短時(shí)間訓(xùn)練當(dāng)今世界最大最前沿的 AI 模型，仍離不開高效的分布式并行擴(kuò)展。針對現(xiàn)有方案并行維度有限、效率不高、通用性差、部署困難、缺乏維護(hù)等痛點(diǎn)，Colossal-AI 通過高效多維并行和異構(gòu)并行等技術(shù)，讓用戶僅需極少量修改，即可高效快速部署 AI 大模型訓(xùn)練。例如對于 GPT-3 這樣的超大 AI 模型，相比英偉達(dá)方案，Colossal-AI 僅需一半的計(jì)算資源，即可啟動(dòng)訓(xùn)練；若使用相同計(jì)算資源，則能提速 11%，可降低 GPT-3 訓(xùn)練成本超百萬美元。對于蛋白質(zhì)結(jié)構(gòu)預(yù)測應(yīng)用 AlphaFold，基于 Colossal-AI 的加速方案的 FastFold，成功超越谷歌和哥倫比亞大學(xué)的方案，將 AlphaFold 訓(xùn)練時(shí)間從 11 天減少到 67 小時(shí)，且總成本更低，在長序列推理中也實(shí)現(xiàn) 9.3～11.6 倍的速度提升。Colossal-AI 注重開源社區(qū)建設(shè)，提供中文教程，開放用戶社群及論壇，對于用戶反饋進(jìn)行高效交流與迭代更新，不斷添加 PaLM、AlphaFold 等前沿應(yīng)用。自開源以來，Colossal-AI 已經(jīng)多次登頂 GitHub 熱榜 Python 方向世界第一，與眾多已有數(shù)萬 star 的明星開源項(xiàng)目一起受到海內(nèi)外關(guān)注！潞晨技術(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 超算杰出新人獎(jiǎng)、UC 伯克利 EECS Lotfi A. Zadeh 優(yōu)秀畢業(yè)生獎(jiǎng)潞晨 CSO Prof. James Demmel：加州大學(xué)伯克利分校杰出教授、ACM/IEEE Fellow，美國科學(xué)院、工程院、藝術(shù)與科學(xué)院三院院士GitHub 地址：https://github.com/hpcaitech/ColossalAIhttps://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc
本站僅提供存儲(chǔ)服務(wù)，所有內(nèi)容均由用戶發(fā)布，如發(fā)現(xiàn)有害或侵權(quán)內(nèi)容，請
點(diǎn)擊舉報(bào)。
