當涉及訓練大型 AI 模型時，人們會想到使用數千個 GPU，昂貴的訓練成本，只有少數幾家科技巨頭能夠承受。而 AI 用戶，如初創公司或大學的研究人員，只能眼睁睁地看著大型模型的訓練，別無他法。
現在，一台只有一個 GPU 的個人電腦就可以訓練高達 180 億參數的 GPT，一台筆記本電腦也可以訓練具有超過 10 億個參數的模型。與現有的主流解決方案相比，參數容量可以增加十倍以上！
這樣顯著的改進來自於 Colossal-AI，它是一個通用的大型 AI 模型高效訓練系統。最重要的是，它完全開源，只需要進行最小的修改即可讓現有的深度學習項目在單個消費級圖形卡上訓練更大的模型，讓每個人都可以在家中訓練大型 AI 模型！特別是，它使下游任務和應用部署，如大型 AI 模型微調和推理變得更加容易！
通過提供各種受歡迎的高效並行處理，Colossal-AI 還可以幫助其用戶輕鬆地將現有項目部署到大型計算集群中。
自 2018 年谷歌提出具有 3 億參數的 BERT 模型以來，僅僅幾年時間，大型模型的參數記錄就多次被更新，如 OpenAI 提出的具有 1750 億參數的 GPT-3，由微軟和 NVIDIA 共同推出的具有 5300 億參數的 MT-NLG 等等...
密集模型已經實現了超過千億參數的規模，而稀疏的專家混合模型（MoE），例如谷歌於 2021 年發布的 Switch Transformer，已將參數數量提升到萬億級別。
然而，從頭開始訓練如此龐大的模型可能會非常昂貴。通常需要同時使用數百甚至數千個專業高性能 GPU，如 NVIDIA A100。如果我們使用專用的 InfiniBand 高速網絡來構建超級計算機集群，訓練成本甚至可能達到 1000 萬美元。
一塊消費級 GPU，即可訓練大模型#
很明顯，像大學生和個人開發者這樣的 AI 用戶無法承擔這樣高昂的成本來訓練大模型，而在 AI 社區中，這類人最流行的計算資源就是 NVIDIA RTX GPU。
為了提高 AI 生產力，讓更多的開發者受益於大模型，並真正實現我們的願景，使大型 AI 模型的使用 “快速且廉價”，Colossal-AI 只需要幾行代碼就可以實現模型訓練容量的增加十倍。這將使每個人都能夠使用單個普通 GPU 來訓練大型 AI 模型。
在所有類型的硬件上，Colossal-AI 的表現都比普通的 PyTorch 和主流的分佈式解決方案，如微軟的 DeepSpeed 更好。
對於大型模型的代表 ——GPT，Colossal-AI 能夠在一台搭載 RTX 2060 6GB 的遊戲筆記本上訓練最多 1.5 億個參數的 GPT。對於搭載 RTX3090 24GB 的電腦，Colossal-AI 可以幫助用戶訓練具有 180 億個參數的 GPT。對於 Tesla V100 等高性能顯卡，Colossal-AI 也可以帶來顯著的改進。
Colossal-AI 還成功地實現了谷歌最近發布的 PaLM（Pathways 語言模型）。在各種硬件上，它也表現出了出色的性能提升，而微軟 DeepSpeed 尚未發布其官方的 PaLM 實現。
關鍵技術：增強異構訓練#
使用單個消費級 GPU 來訓練大型 AI 模型的最大問題是 GPU 內存容量極其有限，這嚴重限制了可以容納的模型參數。微軟 DeepSpeed 提出的 ZeRO-offload 方法嘗試分割模型並利用 CPU 內存，其容量更大且成本更低。目前，已經有幾個基於 DeepSpeed 進行的異構訓練的修改版本。但是，如下圖左側所示，當 GPU 內存不足以滿足其相應的模型需求時，即使 CPU 上仍有可用內存，系統也會崩潰。
Colossal-AI 與 DeepSpeed 派生的方法不同，它們從頭開始構建了自己的核心技術，如 ZeRO，解決了 DeepSpeed 存在的問題。DeepSpeed 只在 CPU 和 GPU 內存之間靜態地分割模型數據，並且對不同的訓練配置只使用固定的內存佈局。Colossal-AI 做了很多改進，提高了 GPU 和 CPU 內存的使用效率。畢竟，CPU 內存比帶有大內存的高性能顯卡便宜得多。
Colossal-AI 設計的 Gemini 機制有效地管理和利用 GPU 和 CPU 的異構內存，因此在訓練過程中，張量會動態地分配到 CPU-GPU 存儲空間中，從而突破了 GPU 的內存限制。
我們利用深度學習網絡訓練過程的迭代性質，將訓練分為兩個階段：預熱階段和非預熱階段，根據迭代次數劃分。在初始的預熱階段，監視內存信息；在非預熱階段，使用收集到的信息來高效地移動張量，以最小化 CPU-GPU 數據移動。
這聽起來很簡單，但實現起來卻很麻煩。因為非模型數據的生命周期不由用戶管理，而現有的深度學習框架也沒有向用戶公開非模型數據的跟蹤接口，因此很難獲得非模型數據的內存使用情況。其次，還需要考慮非框架開銷，如 CUDA 上下文等。
Colossal-AI 通過在熱身階段進行採樣來獲取 CPU 和 GPU 內存的使用情況。而非模型數據的使用情況可以通過比較兩個時刻之間的系統最大內存使用量和模型內存使用量來獲得。通過查詢內存管理器，可以知道模型的內存使用情況，如下圖所示的黑色實線所示。
這聽起來很簡單，但實現起來卻很麻煩。因為非模型數據的生命周期不由用戶管理，而現有的深度學習框架也沒有向用戶公開非模型數據的跟蹤接口，因此很難獲得非模型數據的內存使用情況。其次，還需要考慮非框架開銷，如 CUDA 上下文等。
Colossal-AI 通過在熱身階段進行採樣來獲取 CPU 和 GPU 內存的使用情況。而非模型數據的使用情況可以通過比較兩個時刻之間的系統最大內存使用量和模型內存使用量來獲得。通過查詢內存管理器，可以知道模型的內存使用情況，如下圖所示的黑色實線所示。
所有模型張量都由內存管理器管理，每個張量都標有狀態信息，包括 HOLD、COMPUTE、FREE 等。根據動態查詢的內存使用情況，Colossal-AI 改變張量狀態，並不斷調整張量位置。最終實現了 GPU 和 CPU 內存的高效使用，最大化了模型容量，並在極其有限的硬件情況下平衡了訓練速度。這對於 AI 民主化和大型模型的低成本微調的下游任務具有重要意義。
此外，分佈式並行技術是進一步加速模型訓練的重要方法。針對現有解決方案的痛點，如並行維度有限、效率低、通用性差、部署困難和缺乏維護，Colossal-AI 使用了多維並行和異構並行等技術，使用戶可以僅通過對其代碼進行少量修改，即可高效快速地部署大型 AI 模型。
例如，對於像 GPT-3 這樣的超大型 AI 模型，與 NVIDIA 解決方案相比，Colossal-AI 只需要一半的計算資源即可開始訓練；如果使用相同的計算資源，速度可以進一步提高 11％，從而可以將 GPT-3 的訓練成本降低超過 100 萬美元。
對於用於蛋白質結構預測的 AlphaFold，我們團隊引入了基於 Colossal-AI 加速方案的 FastFold。FastFold 成功超越了谷歌和哥倫比亞大學提出的其他方案，將 AlphaFold 的訓練時間從 11 天縮短至 67 小時，同時也降低了總成本。此外，在長序列推理方面，我們實現了 9.3~11.6 倍的速度提升。
此外，Colossal-AI 非常重視開源社區建設，提供英文和中文教程，並支持最新的前沿應用，如 PaLM 和 AlphaFold。我們還將定期推出新的創新功能。我們始終歡迎社區的建議和討論，如果您遇到任何問題，我們將非常樂意為您提供幫助。您可以在這裡提出問題或在我們的論壇中創建討論主題。我們非常感謝您的建議。最近，Colossal-AI 在 Github 的熱門項目排行榜中排名第一，而這是在許多擁有 10K 星的項目的背景下實現的。
關於原作者
原作者都是 HPC-AI Tech 的核心成員，來自加州大學伯克利分校、斯坦福大學、清華大學、北京大學、新加坡國立大學、南洋理工大學等知名高校，同時他們也曾在 Google Brain、IBM、英特爾、微軟、NVIDIA 等科技巨頭工作過。該公司也獲得了創新工場、真格基金等多家頂級風險投資機構的種子輪融資。
Prof. Yang You, Founder of HPC-AI Tech
Ph.D., University of California, Berkeley
IPDPS/ICPP Best Paper Author
ACM/IEEE CS George Michael Memorial HPC Fellowship
Forbes 30 Under 30 (Asia 2021)
IEEE-CS Outstanding Newcomer Award in Supercomputing
UC Berkeley EECS Lotfi A. Zadeh Prize
Prof. James Demmel, CSO of HPC-AI Tech
Distinguished Professor, University of California, Berkeley
ACM/IEEE Fellow
Member of the American Academy of Sciences, the Academy of Engineering, and the Academy of Arts and Sciences
Check out the project over here: https://github.com/hpcaitech/ColossalAI
原文地址（本文對原文進行了翻譯）：https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc