Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

ADVERTISEMENT

首個開源的ChatGPT低成本重現流程來了! 

預訓練、獎勵模型訓練、強化學習訓練,一次性打通。 

最小demo訓練流程僅需1.62GB視訊記憶體,隨便一張消費級顯示卡都能滿足了。單卡模型容量最多提升10.3倍。 

相比原生PyTorch,單機訓練速度最高可提升7.73倍,單卡推理速度提升1.42倍,僅需一行程式碼即可呼叫。對於微調任務,可最多提升單卡的微調模型容量3.7倍,同時保持高速運行,同樣僅需一行程式碼。 

要知道,ChatGPT紅是真的紅,但要重現也是真的難。 

畢竟ChatGPT是不開源的,市面上至今沒有開源預訓練權重、完全開源的低成本訓練流程,而且千億等級大模型的訓練本身就是個難題。 

但ChatGPT軍備賽已經愈演愈烈,為了抓住趨勢,如Google等都在打造對抗的產品。快速重現ChatGPT是應趨勢所需。 

開源加速方案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版本,以降低硬體限制。

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

要知道,在單機多卡伺服器上,即便把顯示卡規格提升為A100 80GB,由於ChatGPT的複雜性和記憶體碎片,PyTorch最大也只能啟動基於GPT-L(774M)這樣的小模型ChatGPT。 

用PyTorch原生的DistributedDataParallel (DDP) 進行多卡並行擴充至4卡或8卡,性能提升有限。 

Colossal-AI最高可提升單機訓練速度7.73倍,單卡推理速度1.42倍,還可繼續擴大規模並行。 

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

為了盡可能降低訓練成本和上手門檻,Colossal-AI還提供了在單張GPU上即可嘗試的ChatGPT訓練流程。 

相比於PyTorch在約50萬元的A100 80GB上,最大僅能啟動7.8億參數模型,Colossal-AI將單卡容量提升10.3倍至80億參數。 

對於基於1.2億參數小模型的ChatGPT訓練,最低僅需1.62GB視訊記憶體,任意單張消費級GPU即可滿足。 

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

此外,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作為系統原則即可快速使用。

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等): 

# 使用單機單卡訓練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_

背後原理如何?

背後原理如何?

核心方案還是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方法認為大語言模型是過參數化的,其在微調中的參數改變量是一個低秩的矩陣,可以將其分解為兩個更小的的矩陣的乘積,即

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源 

在微調時,固定大模型參數,只調整低秩矩陣參數,從而顯著減小訓練參數量。在微調之後,進行推理部署之前,只需要將參數加回原有矩陣即可,即

Colossal-AI用單張消費級顯示卡,就可以低成本重現ChatGPT完整演算流程,程式現在已開源

 

,不增加模型的推理延遲。

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有好的想法,可以和他們聯繫。

 

 

開源地址:

資料來源:

相關文章: ChatGPT歷史溯源:從1950年代講起,談聊天機器人的前世今生

 

Qbitai
作者

量子位(Qbitai)專注於人工智慧及前沿科技領域,提供技術研發趨勢、科技企業動態、新創公司報道等最新資訊,以及機器學習入門資源、電腦科學最新研究論文、開源程式碼和工具的相關報導。

使用 Facebook 留言
發表回應
謹慎發言,尊重彼此。按此展開留言規則