ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

最新AI资讯1年前 (2023)发布 Aibot114
403 0
广告也精彩
ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

首個開源 ChatGPT 低成本再現流程來了!預訓練、獎勵模型訓練、強化學習訓練,一次性打通。

明敏發自凹非寺,量子位 | 公眾號QbitAI

最小demo訓練流程僅需1.62GB顯卡記憶體,隨便一張消費級顯卡就能滿足。單卡模型容量最多提升10.3倍,相比原生PyTorch,單機訓練速度最高可提升7.73倍,單卡推理速度提升1.42倍,僅需一行代碼。

微調任務可最多提升單卡微調模型容量3.7倍,同時保持高速運行,同樣僅需一行代碼。ChatGPT紅是紅,再現也真的難,畢竟ChatGPT不開源,市面至今沒有開源預訓練權重、完全開源的低成本訓練流程,且千億等級大模型訓練本身就是難題。

但ChatGPT軍備賽已愈演愈烈,為了抓住趨勢,如Google等都在打造競品。快速再現ChatGPT是應趨勢所需,開源加速方案Colossal-AI正是為此而來,且提供開源完整再現流程同時,也降低成本。

降低顯卡記憶體消耗是關鍵

ChatGPT效果好,主要是訓練過程引入人類回饋強化學習(RLHF),但也直接導致ChatGPT再現訓練難度飆升。流程分為三階段:

1. 監督微調:從Prompt庫採樣,收集人工回答,利用這些數據微調預訓練大語言模型。

2. 獎勵模型:從Prompt庫採樣,使用大語言模型產生多個回答,人工排序回答後,訓練獎勵模型(RM)模擬人類價值判斷。

3. 基於第一階段的監督微調模型和第二階段獎勵模型,利用強化學習演算法訓練大語言模型。

ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

▲ RLHF三個階段。

對ChatGPT訓練而言,第三階段是核心部分。OpenAI採用強化學習近端策略最佳化演算法(PPO),引入獎勵訊號,使語言模型產生內容更符合人類評判標準。但引入強化學習,代表有更多模型調用。如使用基於Actor-Critic(AC)結構的PPO演算法,需訓練時進行Actor、Critic兩個模型前向推理和反向傳播,以及監督微調模型、獎勵模型多次前向推理。

ChatGPT基礎的InstructGPT論文,Actor和監督微調模型都使用1,750億參數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能用原本一半硬體資源啟動1,750億參數模型訓練,從64卡降低到32卡。如果續用64卡,則將訓練時間壓縮到更短,節省訓練成本、加速產品更新。

為了讓更多開發者體驗再現ChatGPT,除了1750億參數版本,Colossal-AI還提供單卡、單機4 / 8卡類ChatGPT版本,以降低硬體限制。ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

單機多卡服務器,即便顯卡規格提升為A100 80GB,由於ChatGPT複雜性和內存碎片,PyTorch最大也只能啟動基於GPT-L(774M)這類小模型ChatGPT。用PyTorch原生DistributedDataParallel(DDP)進行多卡並行擴展至4卡或8卡,性能提升有限。Colossal-AI最高可提升單機訓練速度7.73倍,單卡推理速度1.42倍,還可繼續擴大規模並行。

ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

為了盡可能降低訓練成本和上手門檻,Colossal-AI還提供單張GPU即可嘗試的ChatGPT訓練流程。相比PyTorch在約10萬元的A100 80GB最大僅能啟動7.8億參數模型,Colossal-AI將單卡容量提升10.3倍至80億參數。對於基於1.2億參數小模型的ChatGPT訓練,最低僅需1.62GB顯卡記憶體,任意單張消費級GPU即可滿足。

ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

此外,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, 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等):

# 使用單機單卡訓練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超過8千,並成功入選SC、AAAI、PPoPP、CVPR等國際AI與HPC頂級會議的官方教程。

目前Colossal-AI已成功幫某世界500強企業開發具增強線上搜尋引擎能力的類ChatGPT聊天機器人模型,之前也為Stable Diffusion、OPT、AlphaFold等先進模型提供多樣高效大規模多維並行分布式解決方案。主創者為加州柏克萊大學教授James Demmel和新加坡國立大學校長尤洋。

ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

▲ Colossal-AI與主要開源專案同期開源數據對比。

細節原理方面,LoRA、ZeRO+Gemini是關鍵。

低成本微調的LoRA

微調部分,Colossal-AI支援使用低秩矩陣微調(LoRA)法。LoRA法認為大語言模型過參數化,微調參數改變量是低秩矩陣,可分解為兩個更小矩陣的乘積:ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

微調時固定大模型參數,只調整低秩矩陣參數,減少訓練參數量。微調後推理部署前,只需將參數加回原有矩陣即可:ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

不增加模型的推理延遲:ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

▲ LoRA示意圖,僅需訓練A、B。

減少記憶體冗餘的ZeRO+Gemini

Colossal-AI支援使用無冗餘優化器(ZeRO)最佳化記憶體使用,可有效減少記憶體冗餘,並比傳統數據並行策略更不犧牲計算粒度和通訊效率,同時大幅提高記憶體使用效率。

為了提升ZeRO性能,Colossal-AI引入自動Chunk機制。透過將運算順序連續一組參數存入同Chunk(Chunk是一段連續記憶體空間),確保每個Chunk大小相同,提高記憶體使用效率。使用Chunk法組織記憶體可確保PCI-e和GPU-GPU之間網路頻寬有效利用,減少通訊次數,同時避免潛在記憶體碎片。ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

▲ Chunk機制。

此外,Colossal-AI異構記憶體空間管理器Gemini支援將優化器狀態從GPU卸除到CPU ,以節省GPU記憶體。同時利用GPU、CPU記憶體(由CPU DRAM或NVMe SSD記憶體組成)突破單GPU記憶體牆限制,擴展可訓練模型規模。ChatGPT 低成本再現流程開源!任意單張消費級顯卡可體驗,記憶體需求低至 1.62GB

▲ 透過ZeRO+Gemini提升硬體模型容量。

One More Thing

儘管這次開源包含再現ChatGPT的完整演算法流程和必要軟體系統,但想要實際應用,至少還需數據、演算力等支援。參考開源大模型BLOOM、開源AI畫畫工具Stable Diffusion的經驗,背後都需個人開發者、演算力、數據模型等可能合作方支援共建。

以前超過千名科學家聯合發起、耗時一年多煉成的號稱和GPT-3一樣強大的語言模型BLOOM,還有AI繪圖頭號明星Stable Diffusion,也是由Stability AI、EleutherAI和LAION聯合完成。再現ChatGPT也是如此,Colossal-AI正在發起活動。如果讀者對這工作感興趣或有好想法,可透過以下方式與他們聯絡:

GitHub發issue或提交PR。加入Colossal-AI用戶微信或Slack群。點擊「閱讀原文」填寫合作提案。發送合作提案到信箱[email protected]。Open source solution replicates ChatGPT training process! Ready to go with only 1.6GB GPU memory and gives you 7.73 times faster training!

(作者:明敏;本文由 量子位 QbitAI 授權轉載;首圖來源:shutterstock)

© 版权声明

相关文章

暂无评论

暂无评论...