火爆全網的 ChatGPT，彷彿開啟了第四次工業革命，讓微軟、Google等全球科技巨頭打得昏天黑地，引得各路玩家紛紛入局，搶佔賽道。
然而由於 OpenAI 沒有開源 ChatGPT，如何有效復現 ChatGPT 已成為擺在大家面前的頭號難題，急需可靠的開源共建方案。
Colossal-AI 快速跟進，首個開源低成本復現 ChatGPT 完整流程。作為當下最火熱的開源 AI 大模型解決方案，Colossal-AI 已收穫開源社區GitHub Star 近萬顆，此次開源亮點包括：
- 開源完整基於 PyTorch 的 ChatGPT 復現流程，涵蓋全部 3 個階段，可實現從預訓練模型到 ChatGPT 的蛻變；
- 體驗最小 demo 訓練流程最低僅需 1.62GB 視訊記憶體，任意單張消費級 GPU 即可滿足，單卡模型容量最多提升 10.3 倍；
- 相比原生 PyTorch，最高可提升單機訓練速度 7.73 倍，單卡推理速度 1.42 倍，一行程式碼即可使用；
- 對於微調任務，可最多提升單卡的微調模型容量 3.7 倍，同時保持高速運行，僅需一行程式碼；
- 提供單卡、單機多卡、1750 億參數等多個版本，支持從 Hugging Face 匯入 OPT，GPT-3，BLOOM 等多種預訓練大模型；
- 收斂驗證正在進行中，該項目也在吸引合作者共建生態。
開源地址：https://github.com/hpcaitech/ColossalAI
ChatGPT——AIGC 引發的工業革命
如果問新年伊始，最火爆的科技熱點是什麼？非 ChatGPT 莫屬。
它彷彿無所不能的六邊形戰士，可以聊天、寫程式碼、修改 bug、做表格、發論文、寫作業、做翻譯、甚至代替 Google 搜尋引擎等……
自發布以來，ChatGPT 便已摧枯拉朽之勢席捲各個行業，不僅 5 天時間便突破百萬使用者，月活使用者突破 1 億更是僅用時 2 個月，成為史上增速最快的消費級應用，遠超如今其他知名應用，如 Twitter 5 年、Meta（Facebook）4 年半，TikTok 9 個月等，而手機普及到 1 億使用者則用了 16 年。
1 億使用者月活使用者耗時
比爾・蓋茲盛讚「ChatGPT 的意義不亞於 PC 和網際網路誕生」，而微軟 CEO 薩蒂亞・納德拉（Satya Nadella）更是直言「堪比工業革命，這輩子第一次見這麼大的技術浪潮」和「AI 正在重塑網際網路」。作為向 OpenAI 投資上百億美元的大金主，微軟已火速將 ChatGPT 整合進自家的搜尋引擎必應 Bing 和 Edge 瀏覽器，還計劃加入 Teams 以及 Office 等辦公套件全家桶，股價一夜市值飆漲超 800 億美元。
微軟與Google發佈會後股價對比
而隔壁需要擔心被 ChatGPT 革命掉自家搜尋引擎的Google，雖然拉響 「紅色警報」，緊急發佈對標競品 Bard，卻因 Demo 首秀翻車，股價市值瞬間蒸發 1000 億美元。
一夜之間，全球的科技巨頭們彷彿都回到了自己年輕時的樣子，紛紛宣佈要打造自己的 ChatGPT。
但 ChatGPT 發佈已有數月，市面上不僅沒有預訓練權重開源，連可靠的完整開源訓練流程都仍是空白，更無法實現基於千億大模型的 ChatGPT 全流程高效搭建和應用。臨時上線，號稱 「對標 ChatGPT」 的一眾新品們，因為閉源也難辨真偽。
為什麼 ChatGPT 有如此魔力？復現它又有哪些難點？
ChatGPT 技術分析
ChatGPT 的驚人效果，重要特徵是在訓練過程引入人類反饋強化學習（RLHF），使得模型表現更符合人類價值觀。
ChatGPT 的訓練流程主要分為三個階段：
1. 從 Prompt 庫中取樣，收集其人工回答，利用這些資料來微調預訓練大語言模型。
2. 從 Prompt 庫中取樣，使用大語言模型生成多個回答，人工對這些回答進行排序後，訓練獎勵模型（RM），來擬合人類的價值判斷。
3. 基於階段 1 的監督微調模型和階段 2 的獎勵模型，利用強化學習演算法對大語言模型進一步訓練。
其中階段 3 是 RLHF 訓練的核心部分，OpenAI 採用了強化學習中的近端策略最佳化演算法（PPO），藉此引入獎勵信號，使得語言模型生成內容更加符合人類評判標準。
RLHF 的三個階段
ChatGPT 模型的複雜性在於強化學習的引入會帶來更多模型的調用。例如，使用基於 Actor-Critic（AC）結構的 PPO 演算法，需要在訓練時進行 Actor、Critic 兩個模型的前向推理和反向傳播，以及監督微調模型、獎勵模型的多次前向推理。在 ChatGPT 基礎的 InstructGPT 的論文中，Actor 和監督微調模型都使用了 1750 億參數的 GPT-3 系列模型，Critic 和獎勵模型則使用了 60 億參數的 GPT-3 系列模型。
對於如此多的模型參數，想要啟動原始 ChatGPT 訓練流程，需要數千 GB 的視訊記憶體開銷，顯然遠超單張 GPU 的容納能力，常見的資料並行技術也無能為力。但即使引入張量並行、流水並行對參數進行劃分，也仍需至少 64 張 80GB 的 A100 作為硬體基礎。並且，其中的流水並行由於 bubble 和排程複雜，效率受限，不適合 AIGC 的生成式任務。階段 3 涉及 4 個模型的複雜強化學習訓練流程，進一步給 ChatGPT 的程式碼復現帶來了困難和挑戰。
使用 Colossal-AI 低成本復現 ChatGPT
Colossal-AI 以開源方式復現了 ChatGPT 訓練的基本流程，包括階段 1 預訓練，階段 2 的獎勵模型的訓練，以及最為複雜的階段 3 的強化學習訓練等。
同時，Colossal-AI 通過 ZeRO，Gemini, Chunk-based 記憶體管理等技術，極大地降低 ChatGPT 訓練的視訊記憶體開銷，僅需一半硬體資源即可啟動 1750 億參數模型訓練（64 卡 ->32 卡），顯著降低應用成本。若使用上述相同硬體資源，Colossal-AI 則能以更短時間進行訓練，節省訓練成本，加速產品迭代。
為了讓更多開發者體驗復現 ChatGPT 模型，除 1750 億參數版本外，Colossal-AI 還提供高效的單卡、單機 4/8 卡的類 ChatGPT 版本，以降低硬體限制。
在單機多卡伺服器上，即便使用最高端的 A100 80GB 顯示卡，由於 ChatGPT 的複雜性和記憶體碎片，PyTorch 最大僅能啟動基於 GPT-L（774M）這樣的小模型的 ChatGPT。用 PyTorch 原生的 DistributedDataParallel (DDP) 進行多卡並行擴展至 4 卡或 8 卡，性能提升有限。
Colossal-AI 不僅在單卡速度上訓練和推理優勢明顯，隨著並行規模擴大還可進一步提升，最高可提升單機訓練速度 7.73 倍，單卡推理速度 1.42 倍，還可繼續擴展至大規模並行，顯著降低 ChatGPT 復現成本。
為了儘可能降低訓練成本和上手門檻，Colossal-AI 還提供了在單張 GPU 上即可嘗試的 ChatGPT 訓練流程。相比於 PyTorch 在約 10 萬元的 A100 80GB 上，最大僅能啟動 7.8 億參數模型，Colossal-AI 將單卡容量提升 10.3 倍至 80 億參數。對於基於 1.2 億參數小模型的 ChatGPT 訓練，最低僅需 1.62GB 視訊記憶體，任意單張消費級 GPU 即可滿足。
此外，Colossal-AI 也致力於降低基於預訓練大模型的微調任務成本。以 ChatGPT 可選的開源基礎模型 OPT 為例，相比 PyTorch，Colossal-AI 可將提升單卡微調模型容量 3.7 倍（原始計算量顯著增大），同時保持高速運行。
一行程式碼快速上手
Colossal-AI 為 Hugging Face 社區的 GPT，OPT 和 BLOOM 等主流預訓練模型，提供了開箱即用的 ChatGPT 復現程式碼。以 GPT 為例，僅需一行程式碼，指定使用 Colossal-AI 作為系統策略即可快速使用。
Pythonfrom chatgpt.nn import GPTActor, GPTCritic, RewardModelfrom chatgpt.trainer import PPOTrainerfrom chatgpt.trainer.strategies import ColossalAIStrategystrategy = ColossalAIStrategy (stage=3, placement_policy='cuda')with strategy.model_init_context ():actor = GPTActor ().cuda ()critic = GPTCritic ().cuda ()initial_model = deepcopy (actor).cuda ()reward_model = RewardModel (deepcopy (critic.model)).cuda ()trainer = PPOTrainer (strategy, actor, critic, reward_model, initial_model, ...)trainer.fit (prompts)
使用下列命令，即可快速啟動單卡、單機多卡、1750 億版本訓練，並測試各種性能指標（包括最大視訊記憶體佔用、吞吐率和 TFLOPS 等）：
Python# 使用單機單卡訓練 GPT2-S，使用最小的 batch size，Colossal-AI Gemini CPU 策略torchrun --standalone --nproc_pero_node 1 benchmark_gpt_dummy.py --model s --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1# 使用單機 4 卡訓練 GPT2-XL，使用 Colossal-AI Zero2 策略torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2# 使用 4 機 32 卡訓練 GPT-3，使用 Colossal-AI Gemini CPU 策略torchrun --nnodes 4 --nproc_per_node 8 \--rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR \benchmark_gpt_dummy.py --model 175b --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1
背後最佳化
核心系統 Colossal-AI
復現 ChatGPT 的背後，依賴面向大模型時代的通用深度學習系統 Colossal-AI，可基於 PyTorch 高效快速部署 AI 大模型訓練和推理，降低 AI 大模型應用成本。
自開源以來，Colossal-AI 已經多次在 GitHub 熱榜位列世界第一，獲得 GitHub Star 超八千顆，併成功入選 SC、AAAI、PPoPP、CVPR 等國際 AI 與 HPC 頂級會議的官方教程。除上述最佳化外，Colossal-AI 還針對 AI 大模型趨勢，提供最多樣和高效的大規模多維並行分散式解決方案，此前已在 Stable Diffusion、OPT、AlphaFold 等前沿模型上展現卓越優勢。
Colossal-AI 與當今主要開源項目同期開源資料對比
Colossal-AI 由加州柏克萊大學傑出教授 James Demmel 和新加坡國立大學校長青年教授尤洋領導。相關解決方案已成功在自動駕駛、雲端運算、零售、醫藥、晶片等行業知名廠商落地應用，廣受好評。Colossal-AI 已成功幫助某世界 500 強企業，開發具備線上搜尋引擎能力增強的類 ChatGPT 聊天機器人模型。
低成本微調的 LoRA
Colossal-AI 支持使用低秩矩陣微調（LoRA）方法進行高效微調。LoRA 方法認為大語言模型是過參數化的，其在微調中的參數改變數是·一個低秩的矩陣，可以將其分解為兩個更小的的矩陣的乘積，即
。在微調時，固定大模型參數，只調整低秩矩陣參數，從而顯著減小訓練參數量。在微調之後，進行推理部署之前，只需要將參數加回原有矩陣即可，即
，不增加模型的推理延遲。
LoRA 示意圖，僅需訓練 A、B
減少記憶體冗餘的 ZeRO + Gemini
Colossal-AI 支持使用無冗餘最佳化器 (ZeRO) 來最佳化記憶體使用，這種方法可以有效減少記憶體冗餘，並且相比傳統的資料並行策略，不會犧牲計算粒度和通訊效率，同時可以大幅提高記憶體使用效率。為了進一步提升 ZeRO 的性能，Colossal-AI 引入了自動 Chunk 機制。通過將運算順序上連續的一組參數存入同一個 Chunk 中（Chunk 是一段連續的記憶體空間），可以確保每個 Chunk 的大小相同，從而提高記憶體使用效率。使用 Chunk 方式組織記憶體可以保證 PCI-e 和 GPU-GPU 之間的網路頻寬得到有效利用，減小通訊次數，同時避免潛在的記憶體碎片。
Chunk 機制
此外，Colossal-AI 的異構記憶體空間管理器 Gemini 支持將最佳化器狀態從 GPU 卸載到 CPU ，以節省 GPU 記憶體佔用。可以同時利用 GPU 記憶體、CPU 記憶體（由 CPU DRAM 或 NVMe SSD 記憶體組成）來突破單 GPU 記憶體牆的限制，進一步擴展了可訓練模型規模。
通過 ZeRO + Gemini 提升硬體的模型容量
開放協作
儘管此次開源包含了復現 ChatGPT 的完整演算法流程和必要軟體系統，但對於像 ChatGPT 這樣的超大 AI 大模型，想要實際落地應用，還需要資料、算力至少 2 方面的努力。畢竟訓練一個 1750 億參數的 GPT-3 就需要數百萬美元算力。因此，長期以來預訓練大模型都由少數大型私營科技公司壟斷。
好在開源社區已成功進行了新的嘗試。例如，完全開放程式碼、資料集、權重的 1760 億參數的 BLOOM 模型，共有來自全球 60 個國家、超過 250 個機構，以及超過 1000 名研究人員參與其中，其中包括以個人名義參加的 Meta、Google等大廠員工。而前段時間大火的開源圖文生成模型 Stable Diffusion，也是由 Stability AI、EleutherAI 和 LAION 等組織共同完成的。
借鑑上述成功模式，該項目也在吸引更多的合作者：無論是個人開發者，還是算力、資料、模型等可能合作方，都有機會參與其中，大顯身手，以復現 ChatGPT 為起點，擁抱大模型時代！
開源地址：
https://github.com/hpcaitech/ColossalAI
參考連結：
https://www.hpc-ai.tech/blog/colossal-ai-chatgpt
