本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。
首個開源的ChatGPT低成本復現流程來了！
預訓練、獎勵模型訓練、強化學習訓練，一次性打通。
最小demo訓練流程僅需1.62GB顯存，隨便一張消費級顯卡都能滿足了。單卡模型容量最多提升10.3倍。
相比原生PyTorch，單機訓練速度最高可提升7.73倍，單卡推理速度提升1.42倍，僅需一行代碼即可調用。
對于微調任務，可最多提升單卡的微調模型容量3.7倍，同時保持高速運行，同樣僅需一行代碼。
要知道，ChatGPT火是真的火，復現也是真的難。
畢竟ChatGPT是不開源的，市面上至今沒有開源預訓練權重、完全開源的低成本訓練流程，而且千億級別大模型的訓練本身就是個難題。
但ChatGPT軍備賽已經愈演愈烈，為了抓住趨勢，如谷歌等都在打造對標競品?？焖購同FChatGPT是應趨勢所需。
開源加速方案Colossal-AI正是為此而來。
并且在提供開源完整復現流程的同時，把成本降了下來！
開源地址：https://github.com/hpcaitech/ColossalAI
降顯存開銷是關鍵
ChatGPT的效果好，主要是由于在訓練過程中引入了人類反饋強化學習（RLHF），但這也直接導致ChatGPT的復現訓練難度飆升。
其訓練流程主要分為三個階段：
1、監督微調：從Prompt庫中采樣，收集其人工回答，利用這些數據來微調預訓練大語言模型；
2、獎勵模型：從Prompt庫中采樣，使用大語言模型生成多個回答，人工對這些回答進行排序后，訓練獎勵模型（RM），來擬合人類的價值判斷。
3、基于第一階段的監督微調模型和第二階段的獎勵模型，利用強化學習算法對大語言模型進一步訓練。
△RLHF的三個階段
對于ChatGPT訓練而言，第三階段是核心部分。
OpenAI采用了強化學習中近端策略優化算法（PPO），借此引入獎勵信號，使得語言模型生成內容更加符合人類評判標準。
但強化學習的引入，也意味著更多模型調用。
例如，使用基于Actor-Critic（AC）結構的PPO算法，需要在訓練時進行Actor、Critic兩個模型的前向推理和反向傳播，以及監督微調模型、獎勵模型的多次前向推理。
在ChatGPT基礎的InstructGPT論文中，Actor和監督微調模型都使用了1750億參數的GPT-3系列模型，Critic和獎勵模型則使用了60億參數的GPT-3系列模型。
如此大規模的模型參數，意味著想要啟動原始ChatGPT訓練流程，需要數千GB的顯存開銷，單張GPU顯然無法容納，常見的數據并行技術也不能搞定。
即便引入張量并行、流水并行對參數進行劃分，也需要至少64張80GB的A100作為硬件基礎。而且流水并行本身并不適合AIGC的生成式任務，bubble和調度復雜會導致效率受限。
單張消費級顯卡都能體驗
Colossal-AI基于ZeRO，Gemini, LoRA, Chunk-based內存管理等方法，提出了一系列單卡、單機多卡、大規模并行解決方案。
對于基于GPT-3系列模型的ChatGPT，Colossal-AI能用原本一半的硬件資源啟動1750億參數模型訓練，從64卡降低到32卡。
如果繼續用64卡，則將訓練時間壓縮到更短，節省訓練成本、加速產品迭代。
而為了能讓更大范圍的開發者體驗復現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倍，還可繼續擴大規模并行。
為了盡可能降低訓練成本和上手門檻，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開源方案，都能實現一行代碼快速上手。
先看模型使用方面。
盡管ChatGPT背后的大語言模型GPT-3.5不開源，但如GPT、OPT、BLOOM等主流開源模型可作為替代。
Colossal-AI為Hugging Face社區的這些模型，提供了開箱即用的ChatGPT復現代碼，可覆蓋三個階段的訓練。
以GPT為例，添加一行代碼指定使用Colossal-AI作為系統策略即可快速使用。
使用下列命令，即可快速啟動單卡、單機多卡、1750億版本訓練，并測試各種性能指標（包括最大顯存占用、吞吐率和TFLOPS等）：
背后原理如何？
核心方案還是Colossal-AI。
它從誕生起就面向大模型應用，可基于PyTorch高效快速部署AI大模型訓練和推理，是這一領域的明星項目了，GitHub Star超八千顆，并成功入選SC、AAAI、PPoPP、CVPR等國際AI與HPC頂級會議的官方教程。
目前，Colossal-AI已成功幫助一家世界五百強企業，開發具備在線搜索引擎能力增強的類ChatGPT聊天機器人模型。
此前，它們還為Stable Diffusion、OPT、AlphaFold等前沿模型，提供了多樣高效的大規模多維并行分布式解決方案。
主創人員為加州伯克利大學杰出教授James Demmel和新加坡國立大學校長青年教授尤洋。
△Colossal-AI與當今主要開源項目同期開源數據對比
具體到細節原理上，LoRA、ZeRO+Gemini是關鍵。
低成本微調的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提升硬件的模型容量
One More Thing
盡管此次開源包含了復現ChatGPT的完整算法流程和必要軟件系統，但想要走到實際應用落地，還至少需要數據、算力等方面的支持。
參考開源大模型BLOOM、開源AI畫畫工具Stable Diffusion的經驗，這背后都需要包括個人開發者、算力、數據模型等可能合作方的支持共建——
此前，超過1000個科學家聯合發起、耗時一年多煉出了號稱和GPT-3一樣強大的語言模型BLOOM。還有AI畫畫趨勢的頭號明星Stable Diffusion，也是由Stability AI、EleutherAI和LAION多方聯合完成的。
復現ChatGPT也是如此，Colossal-AI正在發起這一開發活動。
如果你對這項工作感興趣or有好的想法，可通過以下方式與他們取得聯系：
- 在GitHub發布issue或提交PR
- 加入Colossal-AI用戶微信或Slack群交流
- 點擊“閱讀原文”填寫合作提案
- 發送合作提案到郵箱contact@hpcaitech.com
傳送門：
開源地址：
?https://github.com/hpcaitech/ColossalAI
