開源方案復現ChatGPT流程！1.62GB視訊記憶體即可體驗，單機訓練提速7.73倍

語言: CN / TW / HK
時間 2023-02-15 16:02:40
主題:

火爆全網的ChatGPT，彷彿開啟了第四次工業革命，讓微軟、谷歌等全球科技巨頭打得昏天黑地，引得各路玩家紛紛入局，搶佔賽道。

然而由於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等多種預訓練大模型；
  • 收斂驗證正在進行中，該專案也在吸引合作者共建生態。

開源地址：http://github.com/hpcaitech/ColossalAI

 

ChatGPT——AIGC引發的工業革命

如果問新年伊始，最火爆的科技熱點是什麼？非ChatGPT莫屬。

它彷彿無所不能的六邊形戰士，可以聊天、寫程式碼、修改 bug、做表格、發論文、寫作業、做翻譯、甚至代替Google搜尋引擎等……

自發布以來，ChatGPT便已摧枯拉朽之勢席捲各個行業，不僅5天時間便突破百萬使用者，月活使用者突破1億更是僅用時2個月，成為史上增速最快的消費級應用，遠超如今其他知名應用，如Twitter 5 年、Meta（Facebook）4 年半，TikTok 9個月等，而手機普及到1億使用者則用了16年。


               1億使用者月活使用者耗時
                                    1億使用者月活使用者耗時

比爾·蓋茨盛讚“ChatGPT的意義不亞於PC和網際網路誕生”，而微軟CEO薩蒂亞·納德拉（Satya Nadella）更是直言“堪比工業革命，這輩子第一次見這麼大的技術浪潮” 和“AI 正在重塑網際網路”。作為向OpenAI投資上百億美元的大金主，微軟已火速將ChatGPT整合進自家的搜尋引擎必應Bing和Edge瀏覽器，還計劃加入Teams以及Office等辦公套件全家桶，股價一夜市值飆漲超800億美元。


               
                                      微軟與谷歌釋出會後股價對比

而隔壁需要擔心被ChatGPT革命掉自家搜尋引擎的谷歌，雖然拉響“紅色警報”，緊急釋出對標競品Bard，卻因Demo首秀翻車，股價市值瞬間蒸發1000億美元。

一夜之間，全球的科技巨頭們彷彿都回到了自己年輕時的樣子，紛紛宣佈要打造自己的ChatGPT。

但ChatGPT釋出已有數月，市面上不僅沒有預訓練權重開源，連可靠的完整開源訓練流程都仍是空白，更無法實現基於千億大模型的ChatGPT全流程高效搭建和應用。臨時上線，號稱“對標ChatGPT”的一眾新品們，因為閉源也難辨真偽。

為什麼ChatGPT有如此魔力？復現它又有哪些難點？

 

ChatGPT技術分析

ChatGPT的驚人效果，重要特徵是在訓練過程引入人類反饋強化學習（RLHF），使得模型表現更符合人類價值觀。

ChatGPT的訓練流程主要分為三個階段：

  1. 從Prompt庫中取樣，收集其人工回答，利用這些資料來微調預訓練大語言模型。
  1. 從Prompt庫中取樣，使用大語言模型生成多個回答，人工對這些回答進行排序後，訓練獎勵模型（RM），來擬合人類的價值判斷。
  1. 基於階段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作為系統策略即可快速使用。

Python
from chatgpt.nn import GPTActor, GPTCritic, RewardModel
from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import ColossalAIStrategy

strategy = 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、谷歌等大廠員工。而前段時間大火的開源圖文生成模型Stable Diffusion，也是由Stability AI、EleutherAI和LAION等組織共同完成的。

借鑑上述成功模式，該專案也在吸引更多的合作者：無論是個人開發者，還是算力、資料、模型等可能合作方，都有機會參與其中，大顯身手，以復現ChatGPT為起點，擁抱大模型時代！

可通過以下方式聯絡或參與：

  • 在GitHub釋出issue或提交PR
  • 加入Colossal-AI使用者微信或Slack群交流
  • 點選閱讀原文填寫合作提案
  • 傳送合作提案到郵箱[email protected]

 

開源地址：http://github.com/hpcaitech/ColossalAI

 

參考連結：http://www.hpc-ai.tech/blog/colossal-ai-chatgpt

 

（閱讀原文連結）http://www.hpc-ai.tech/partners

 

「」