本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

首個(gè)開(kāi)源的ChatGPT低成本復(fù)現(xiàn)流程來(lái)了！

預(yù)訓(xùn)練、獎(jiǎng)勵(lì)模型訓(xùn)練、強(qiáng)化學(xué)習(xí)訓(xùn)練，一次性打通。

最小demo訓(xùn)練流程僅需1.62GB顯存，隨便一張消費(fèi)級(jí)顯卡都能滿(mǎn)足了。單卡模型容量最多提升10.3倍。

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

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

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

畢竟ChatGPT是不開(kāi)源的，市面上至今沒(méi)有開(kāi)源預(yù)訓(xùn)練權(quán)重、完全開(kāi)源的低成本訓(xùn)練流程，而且千億級(jí)別大模型的訓(xùn)練本身就是個(gè)難題。

但ChatGPT軍備賽已經(jīng)愈演愈烈，為了抓住趨勢(shì)，如谷歌等都在打造對(duì)標(biāo)競(jìng)品?？焖?gòu)?fù)現(xiàn)ChatGPT是應(yīng)趨勢(shì)所需。

開(kāi)源加速方案Colossal-AI正是為此而來(lái)。

并且在提供開(kāi)源完整復(fù)現(xiàn)流程的同時(shí)，把成本降了下來(lái)！

開(kāi)源地址：https://github.com/hpcaitech/ColossalAI

降顯存開(kāi)銷(xiāo)是關(guān)鍵

ChatGPT的效果好，主要是由于在訓(xùn)練過(guò)程中引入了人類(lèi)反饋強(qiáng)化學(xué)習(xí)（RLHF），但這也直接導(dǎo)致ChatGPT的復(fù)現(xiàn)訓(xùn)練難度飆升。

其訓(xùn)練流程主要分為三個(gè)階段：

1、監(jiān)督微調(diào)：從Prompt庫(kù)中采樣，收集其人工回答，利用這些數(shù)據(jù)來(lái)微調(diào)預(yù)訓(xùn)練大語(yǔ)言模型；

2、獎(jiǎng)勵(lì)模型：從Prompt庫(kù)中采樣，使用大語(yǔ)言模型生成多個(gè)回答，人工對(duì)這些回答進(jìn)行排序后，訓(xùn)練獎(jiǎng)勵(lì)模型（RM），來(lái)擬合人類(lèi)的價(jià)值判斷。

3、基于第一階段的監(jiān)督微調(diào)模型和第二階段的獎(jiǎng)勵(lì)模型，利用強(qiáng)化學(xué)習(xí)算法對(duì)大語(yǔ)言模型進(jìn)一步訓(xùn)練。

△RLHF的三個(gè)階段

對(duì)于ChatGPT訓(xùn)練而言，第三階段是核心部分。

OpenAI采用了強(qiáng)化學(xué)習(xí)中近端策略?xún)?yōu)化算法（PPO），借此引入獎(jiǎng)勵(lì)信號(hào)，使得語(yǔ)言模型生成內(nèi)容更加符合人類(lèi)評(píng)判標(biāo)準(zhǔn)。

但強(qiáng)化學(xué)習(xí)的引入，也意味著更多模型調(diào)用。

例如，使用基于Actor-Critic（AC）結(jié)構(gòu)的PPO算法，需要在訓(xùn)練時(shí)進(jìn)行Actor、Critic兩個(gè)模型的前向推理和反向傳播，以及監(jiān)督微調(diào)模型、獎(jiǎng)勵(lì)模型的多次前向推理。

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

如此大規(guī)模的模型參數(shù)，意味著想要啟動(dòng)原始ChatGPT訓(xùn)練流程，需要數(shù)千GB的顯存開(kāi)銷(xiāo)，單張GPU顯然無(wú)法容納，常見(jiàn)的數(shù)據(jù)并行技術(shù)也不能搞定。

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

單張消費(fèi)級(jí)顯卡都能體驗(yàn)

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

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

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

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

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

用PyTorch原生的DistributedDataParallel (DDP) 進(jìn)行多卡并行擴(kuò)展至4卡或8卡，性能提升有限。

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

為了盡可能降低訓(xùn)練成本和上手門(mén)檻，Colossal-AI還提供了在單張GPU上即可嘗試的ChatGPT訓(xùn)練流程。

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

對(duì)于基于1.2億參數(shù)小模型的ChatGPT訓(xùn)練，最低僅需1.62GB顯存，任意單張消費(fèi)級(jí)GPU即可滿(mǎn)足。

此外，Colossal-AI也致力于降低基于預(yù)訓(xùn)練大模型的微調(diào)任務(wù)成本。以ChatGPT可選的開(kāi)源基礎(chǔ)模型OPT為例，相比PyTorch，Colossal-AI可將提升單卡微調(diào)模型容量3.7倍（原始計(jì)算量顯著增大），同時(shí)保持高速運(yùn)行。

一行代碼快速上手

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

先看模型使用方面。

盡管ChatGPT背后的大語(yǔ)言模型GPT-3.5不開(kāi)源，但如GPT、OPT、BLOOM等主流開(kāi)源模型可作為替代。

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

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

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)

使用下列命令，即可快速啟動(dòng)單卡、單機(jī)多卡、1750億版本訓(xùn)練，并測(cè)試各種性能指標(biāo)（包括最大顯存占用、吞吐率和TFLOPS等）：

# 使用單機(jī)單卡訓(xùn)練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
# 使用單機(jī)4卡訓(xùn)練GPT2-XL，使用Colossal-AI Zero2策略
torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2
# 使用4機(jī)32卡訓(xùn)練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。

它從誕生起就面向大模型應(yīng)用，可基于PyTorch高效快速部署AI大模型訓(xùn)練和推理，是這一領(lǐng)域的明星項(xiàng)目了，GitHub Star超八千顆，并成功入選SC、AAAI、PPoPP、CVPR等國(guó)際AI與HPC頂級(jí)會(huì)議的官方教程。

目前，Colossal-AI已成功幫助一家世界五百?gòu)?qiáng)企業(yè)，開(kāi)發(fā)具備在線(xiàn)搜索引擎能力增強(qiáng)的類(lèi)ChatGPT聊天機(jī)器人模型。

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

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

△Colossal-AI與當(dāng)今主要開(kāi)源項(xiàng)目同期開(kāi)源數(shù)據(jù)對(duì)比

具體到細(xì)節(jié)原理上，LoRA、ZeRO+Gemini是關(guān)鍵。

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

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

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

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

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

△LoRA示意圖，僅需訓(xùn)練A、B

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

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

為了進(jìn)一步提升 ZeRO 的性能，Colossal-AI 引入了自動(dòng)Chunk機(jī)制。

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

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

△Chunk機(jī)制

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

可以同時(shí)利用 GPU 內(nèi)存、CPU 內(nèi)存（由 CPU DRAM 或 NVMe SSD內(nèi)存組成）來(lái)突破單GPU內(nèi)存墻的限制，進(jìn)一步擴(kuò)展了可訓(xùn)練模型規(guī)模。

△通過(guò)ZeRO+Gemini提升硬件的模型容量

One More Thing

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

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

此前，超過(guò)1000個(gè)科學(xué)家聯(lián)合發(fā)起、耗時(shí)一年多煉出了號(hào)稱(chēng)和GPT-3一樣強(qiáng)大的語(yǔ)言模型BLOOM。還有AI畫(huà)畫(huà)趨勢(shì)的頭號(hào)明星Stable Diffusion，也是由Stability AI、EleutherAI和LAION多方聯(lián)合完成的。

復(fù)現(xiàn)ChatGPT也是如此，Colossal-AI正在發(fā)起這一開(kāi)發(fā)活動(dòng)。

如果你對(duì)這項(xiàng)工作感興趣or有好的想法，可通過(guò)以下方式與他們?nèi)〉寐?lián)系：

1. 在GitHub發(fā)布issue或提交PR
2. 加入Colossal-AI用戶(hù)微信或Slack群交流
3. 點(diǎn)擊“閱讀原文”填寫(xiě)合作提案
4. 發(fā)送合作提案到郵箱contact@hpcaitech.com

傳送門(mén)：

開(kāi)源地址：
?https://github.com/hpcaitech/ColossalAI