設置
  • 日夜間
    隨系統(tǒng)
    淺色
    深色
  • 主題色

ChatGPT 低成本復現(xiàn)流程開源,任意單張消費級顯卡可體驗,顯存需求低至 1.62 GB

量子位 2023/2/17 14:00:21 責編:夢澤

首個開源的 ChatGPT 低成本復現(xiàn)流程來了!

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

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

相比原生 PyTorch,單機訓練速度最高可提升 7.73 倍,單卡推理速度提升 1.42 倍,僅需一行代碼即可調(diào)用。

對于微調(diào)任務,可最多提升單卡的微調(diào)模型容量 3.7 倍,同時保持高速運行,同樣僅需一行代碼。

要知道,ChatGPT 火是真的火,復現(xiàn)也是真的難。

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

但 ChatGPT 軍備賽已經(jīng)愈演愈烈,為了抓住趨勢,如谷歌等都在打造對標競品??焖購同F(xiàn) ChatGPT 是應趨勢所需。

開源加速方案 Colossal-AI 正是為此而來。

并且在提供開源完整復現(xiàn)流程的同時,把成本降了下來!

開源地址:https://github.com/ hpcaitech / ColossalAI

降顯存開銷是關鍵

ChatGPT 的效果好,主要是由于在訓練過程中引入了人類反饋強化學習(RLHF),但這也直接導致 ChatGPT 的復現(xiàn)訓練難度飆升。

其訓練流程主要分為三個階段:

1、監(jiān)督微調(diào):從 Prompt 庫中采樣,收集其人工回答,利用這些數(shù)據(jù)來微調(diào)預訓練大語言模型;

2、獎勵模型:從 Prompt 庫中采樣,使用大語言模型生成多個回答,人工對這些回答進行排序后,訓練獎勵模型(RM),來擬合人類的價值判斷。

3、基于第一階段的監(jiān)督微調(diào)模型和第二階段的獎勵模型,利用強化學習算法對大語言模型進一步訓練。

RLHF 的三個階段

對于 ChatGPT 訓練而言,第三階段是核心部分。

OpenAI 采用了強化學習中近端策略優(yōu)化算法(PPO),借此引入獎勵信號,使得語言模型生成內(nèi)容更加符合人類評判標準。

但強化學習的引入,也意味著更多模型調(diào)用。

例如,使用基于 Actor-Critic(AC)結構的 PPO 算法,需要在訓練時進行 Actor、Critic 兩個模型的前向推理和反向傳播,以及監(jiān)督微調(diào)模型、獎勵模型的多次前向推理。

在 ChatGPT 基礎的 InstructGPT 論文中,Actor 和監(jiān)督微調(diào)模型都使用了 1750 億參數(shù)的 GPT-3 系列模型,Critic 和獎勵模型則使用了 60 億參數(shù)的 GPT-3 系列模型。

如此大規(guī)模的模型參數(shù),意味著想要啟動原始 ChatGPT 訓練流程,需要數(shù)千 GB 的顯存開銷,單張 GPU 顯然無法容納,常見的數(shù)據(jù)并行技術也不能搞定。

即便引入張量并行、流水并行對參數(shù)進行劃分,也需要至少 64 張 80GB 的 A100 作為硬件基礎。而且流水并行本身并不適合 AIGC 的生成式任務,bubble 和調(diào)度復雜會導致效率受限。

單張消費級顯卡都能體驗

Colossal-AI 基于 ZeRO,Gemini, LoRA, Chunk-based 內(nèi)存管理等方法,提出了一系列單卡、單機多卡、大規(guī)模并行解決方案。

對于基于 GPT-3 系列模型的 ChatGPT,Colossal-AI 能用原本一半的硬件資源啟動 1750 億參數(shù)模型訓練,從 64 卡降低到 32 卡。

如果繼續(xù)用 64 卡,則將訓練時間壓縮到更短,節(jié)省訓練成本、加速產(chǎn)品迭代。

而為了能讓更大范圍的開發(fā)者體驗復現(xiàn) ChatGPT,除了 1750 億參數(shù)版本外,Colossal-AI 還提供單卡、單機 4/8 卡的類 ChatGPT 版本,以降低硬件限制。

要知道,在單機多卡服務器上,即便把顯卡規(guī)格提升為 A100 80GB,由于 ChatGPT 的復雜性和內(nèi)存碎片,PyTorch 最大也只能啟動基于 GPT-L(774M)這樣的小模型 ChatGPT。

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

Colossal-AI 最高可提升單機訓練速度 7.73 倍,單卡推理速度 1.42 倍,還可繼續(xù)擴大規(guī)模并行。

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

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

對于基于 1.2 億參數(shù)小模型的 ChatGPT 訓練,最低僅需 1.62GB 顯存任意單張消費級 GPU 即可滿足。

此外,Colossal-AI 也致力于降低基于預訓練大模型的微調(diào)任務成本。以 ChatGPT 可選的開源基礎模型 OPT 為例,相比 PyTorch,Colossal-AI 可將提升單卡微調(diào)模型容量 3.7 倍(原始計算量顯著增大),同時保持高速運行。

一行代碼快速上手

到了具體操作部分,如上復現(xiàn)流程中的多個步驟,基于 Colossal-AI 開源方案,都能實現(xiàn)一行代碼快速上手。

先看模型使用方面。

盡管 ChatGPT 背后的大語言模型 GPT-3.5 不開源,但如 GPT、OPT、BLOOM 等主流開源模型可作為替代。

Colossal-AI 為 Hugging Face 社區(qū)的這些模型,提供了開箱即用的 ChatGPT 復現(xiàn)代碼,可覆蓋三個階段的訓練。

以 GPT 為例,添加一行代碼指定使用 Colossal-AI 作為系統(tǒng)策略即可快速使用。

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



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 已成功幫助一家世界五百強企業(yè),開發(fā)具備在線搜索引擎能力增強的類 ChatGPT 聊天機器人模型

此前,它們還為 Stable Diffusion、OPT、AlphaFold 等前沿模型,提供了多樣高效的大規(guī)模多維并行分布式解決方案。

主創(chuàng)人員為加州伯克利大學杰出教授 James Demmel 和新加坡國立大學校長青年教授尤洋。

Colossal-AI 與當今主要開源項目同期開源數(shù)據(jù)對比

具體到細節(jié)原理上,LoRA、ZeRO+Gemini 是關鍵。

低成本微調(diào)的 LoRA

在微調(diào)部分,Colossal-AI 支持使用低秩矩陣微調(diào)(LoRA)方法。

LoRA 方法認為大語言模型是過參數(shù)化的,其在微調(diào)中的參數(shù)改變量是一個低秩的矩陣,可以將其分解為兩個更小的的矩陣的乘積,即

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

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

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

減少內(nèi)存冗余的 ZeRO+Gemini

Colossal-AI 支持使用無冗余優(yōu)化器 (ZeRO) 來優(yōu)化內(nèi)存使用,這種方法可以有效減少內(nèi)存冗余,并且相比傳統(tǒng)的數(shù)據(jù)并行策略,不會犧牲計算粒度和通信效率,同時可以大幅提高內(nèi)存使用效率。

為了進一步提升 ZeRO 的性能,Colossal-AI 引入了自動 Chunk 機制。

通過將運算順序上連續(xù)的一組參數(shù)存入同一個 Chunk 中(Chunk 是一段連續(xù)的內(nèi)存空間),可以確保每個 Chunk 的大小相同,從而提高內(nèi)存使用效率。

使用 Chunk 方式組織內(nèi)存可以保證 PCI-e 和 GPU-GPU 之間的網(wǎng)絡帶寬得到有效利用,減小通信次數(shù),同時避免潛在的內(nèi)存碎片。

Chunk 機制

此外,Colossal-AI 的異構內(nèi)存空間管理器 Gemini 支持將優(yōu)化器狀態(tài)從 GPU 卸載到 CPU ,以節(jié)省 GPU 內(nèi)存占用。

可以同時利用 GPU 內(nèi)存、CPU 內(nèi)存(由 CPU DRAM 或 NVMe SSD 內(nèi)存組成)來突破單 GPU 內(nèi)存墻的限制,進一步擴展了可訓練模型規(guī)模。

通過 ZeRO+Gemini 提升硬件的模型容量

One More Thing

盡管此次開源包含了復現(xiàn) ChatGPT 的完整算法流程和必要軟件系統(tǒng),但想要走到實際應用落地,還至少需要數(shù)據(jù)、算力等方面的支持。

參考開源大模型 BLOOM、開源 AI 畫畫工具 Stable Diffusion 的經(jīng)驗,這背后都需要包括個人開發(fā)者、算力、數(shù)據(jù)模型等可能合作方的支持共建 ——

此前,超過 1000 個科學家聯(lián)合發(fā)起、耗時一年多煉出了號稱和 GPT-3 一樣強大的語言模型 BLOOM。還有 AI 畫畫趨勢的頭號明星 Stable Diffusion,也是由 Stability AI、EleutherAI 和 LAION 多方聯(lián)合完成的。

復現(xiàn) ChatGPT 也是如此,Colossal-AI 正在發(fā)起這一開發(fā)活動。

開源地址:

https://github.com/hpcaitech/ColossalAI

參考鏈接:

  • https://www.hpc-ai.tech/blog/colossal-ai-chatgpt

本文來自微信公眾號:量子位 (ID:QbitAI),作者:明敏

廣告聲明:文內(nèi)含有的對外跳轉鏈接(包括不限于超鏈接、二維碼、口令等形式),用于傳遞更多信息,節(jié)省甄選時間,結果僅供參考,IT之家所有文章均包含本聲明。

相關文章

關鍵詞:ChatGPT

軟媒旗下網(wǎng)站: IT之家 最會買 - 返利返現(xiàn)優(yōu)惠券 iPhone之家 Win7之家 Win10之家 Win11之家

軟媒旗下軟件: 軟媒手機APP應用 魔方 最會買 要知