首個開源的 ChatGPT 低成本復(fù)現(xiàn)流程來了！
預(yù)訓(xùn)練、獎勵模型訓(xùn)練、強化學(xué)習(xí)訓(xùn)練，一次性打通。
最小 demo 訓(xùn)練流程僅需 1.62GB 顯存，隨便一張消費級顯卡都能滿足了。單卡模型容量最多提升 10.3 倍。
相比原生 PyTorch，單機訓(xùn)練速度最高可提升 7.73 倍，單卡推理速度提升 1.42 倍，僅需一行代碼即可調(diào)用。
對于微調(diào)任務(wù)，可最多提升單卡的微調(diào)模型容量 3.7 倍，同時保持高速運行，同樣僅需一行代碼。
要知道，ChatGPT 火是真的火，復(fù)現(xiàn)也是真的難。
畢竟 ChatGPT 是不開源的，市面上至今沒有開源預(yù)訓(xùn)練權(quán)重、完全開源的低成本訓(xùn)練流程，而且千億級別大模型的訓(xùn)練本身就是個難題。
但 ChatGPT 軍備賽已經(jīng)愈演愈烈，為了抓住趨勢，如谷歌等都在打造對標(biāo)競品?？焖購?fù)現(xiàn) ChatGPT 是應(yīng)趨勢所需。
開源加速方案 Colossal-AI 正是為此而來。
并且在提供開源完整復(fù)現(xiàn)流程的同時，把成本降了下來！
開源地址：https://github.com/ hpcaitech / ColossalAI
降顯存開銷是關(guān)鍵
ChatGPT 的效果好，主要是由于在訓(xùn)練過程中引入了人類反饋強化學(xué)習(xí)（RLHF），但這也直接導(dǎo)致 ChatGPT 的復(fù)現(xiàn)訓(xùn)練難度飆升。
其訓(xùn)練流程主要分為三個階段：
1、監(jiān)督微調(diào)：從 Prompt 庫中采樣，收集其人工回答，利用這些數(shù)據(jù)來微調(diào)預(yù)訓(xùn)練大語言模型；
2、獎勵模型：從 Prompt 庫中采樣，使用大語言模型生成多個回答，人工對這些回答進(jìn)行排序后，訓(xùn)練獎勵模型（RM），來擬合人類的價值判斷。
3、基于第一階段的監(jiān)督微調(diào)模型和第二階段的獎勵模型，利用強化學(xué)習(xí)算法對大語言模型進(jìn)一步訓(xùn)練。
△RLHF 的三個階段
對于 ChatGPT 訓(xùn)練而言，第三階段是核心部分。
OpenAI 采用了強化學(xué)習(xí)中近端策略優(yōu)化算法（PPO），借此引入獎勵信號，使得語言模型生成內(nèi)容更加符合人類評判標(biāo)準(zhǔn)。
但強化學(xué)習(xí)的引入，也意味著更多模型調(diào)用。
例如，使用基于 Actor-Critic（AC）結(jié)構(gòu)的 PPO 算法，需要在訓(xùn)練時進(jìn)行 Actor、Critic 兩個模型的前向推理和反向傳播，以及監(jiān)督微調(diào)模型、獎勵模型的多次前向推理。
在 ChatGPT 基礎(chǔ)的 InstructGPT 論文中，Actor 和監(jiān)督微調(diào)模型都使用了 1750 億參數(shù)的 GPT-3 系列模型，Critic 和獎勵模型則使用了 60 億參數(shù)的 GPT-3 系列模型。
如此大規(guī)模的模型參數(shù)，意味著想要啟動原始 ChatGPT 訓(xùn)練流程，需要數(shù)千 GB 的顯存開銷，單張 GPU 顯然無法容納，常見的數(shù)據(jù)并行技術(shù)也不能搞定。
即便引入張量并行、流水并行對參數(shù)進(jìn)行劃分，也需要至少 64 張 80GB 的 A100 作為硬件基礎(chǔ)。而且流水并行本身并不適合 AIGC 的生成式任務(wù)，bubble 和調(diào)度復(fù)雜會導(dǎo)致效率受限。
單張消費級顯卡都能體驗
Colossal-AI 基于 ZeRO，Gemini, LoRA, Chunk-based 內(nèi)存管理等方法，提出了一系列單卡、單機多卡、大規(guī)模并行解決方案。
對于基于 GPT-3 系列模型的 ChatGPT，Colossal-AI 能用原本一半的硬件資源啟動 1750 億參數(shù)模型訓(xùn)練，從 64 卡降低到 32 卡。
如果繼續(xù)用 64 卡，則將訓(xùn)練時間壓縮到更短，節(jié)省訓(xùn)練成本、加速產(chǎn)品迭代。
而為了能讓更大范圍的開發(fā)者體驗復(fù)現(xiàn) ChatGPT，除了 1750 億參數(shù)版本外，Colossal-AI 還提供單卡、單機 4/8 卡的類 ChatGPT 版本，以降低硬件限制。
要知道，在單機多卡服務(wù)器上，即便把顯卡規(guī)格提升為 A100 80GB，由于 ChatGPT 的復(fù)雜性和內(nèi)存碎片，PyTorch 最大也只能啟動基于 GPT-L（774M）這樣的小模型 ChatGPT。
用 PyTorch 原生的 DistributedDataParallel (DDP) 進(jìn)行多卡并行擴展至 4 卡或 8 卡，性能提升有限。
Colossal-AI 最高可提升單機訓(xùn)練速度 7.73 倍，單卡推理速度 1.42 倍，還可繼續(xù)擴大規(guī)模并行。
為了盡可能降低訓(xùn)練成本和上手門檻，Colossal-AI 還提供了在單張 GPU 上即可嘗試的 ChatGPT 訓(xùn)練流程。
相比于 PyTorch 在約 10 萬元的 A100 80GB 上，最大僅能啟動 7.8 億參數(shù)模型，Colossal-AI 將單卡容量提升 10.3 倍至 80 億參數(shù)。
對于基于 1.2 億參數(shù)小模型的 ChatGPT 訓(xùn)練，最低僅需 1.62GB 顯存，任意單張消費級 GPU 即可滿足。
此外，Colossal-AI 也致力于降低基于預(yù)訓(xùn)練大模型的微調(diào)任務(wù)成本。以 ChatGPT 可選的開源基礎(chǔ)模型 OPT 為例，相比 PyTorch，Colossal-AI 可將提升單卡微調(diào)模型容量 3.7 倍（原始計算量顯著增大），同時保持高速運行。
一行代碼快速上手
到了具體操作部分，如上復(fù)現(xiàn)流程中的多個步驟，基于 Colossal-AI 開源方案，都能實現(xiàn)一行代碼快速上手。
先看模型使用方面。
盡管 ChatGPT 背后的大語言模型 GPT-3.5 不開源，但如 GPT、OPT、BLOOM 等主流開源模型可作為替代。
Colossal-AI 為 Hugging Face 社區(qū)的這些模型，提供了開箱即用的 ChatGPT 復(fù)現(xiàn)代碼，可覆蓋三個階段的訓(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 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 億版本訓(xùn)練，并測試各種性能指標(biāo)（包括最大顯存占用、吞吐率和 TFLOPS 等）：
# 使用單機單卡訓(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 # 使用單機4卡訓(xùn)練GPT2-XL，使用Colossal-AI Zero2策略 torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2 # 使用4機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)域的明星項目了，GitHub Star 超八千顆，并成功入選 SC、AAAI、PPoPP、CVPR 等國際 AI 與 HPC 頂級會議的官方教程。
目前，Colossal-AI 已成功幫助一家世界五百強企業(yè)，開發(fā)具備在線搜索引擎能力增強的類 ChatGPT 聊天機器人模型。
此前，它們還為 Stable Diffusion、OPT、AlphaFold 等前沿模型，提供了多樣高效的大規(guī)模多維并行分布式解決方案。
主創(chuàng)人員為加州伯克利大學(xué)杰出教授 James Demmel 和新加坡國立大學(xué)校長青年教授尤洋。
具體到細(xì)節(jié)原理上，LoRA、ZeRO+Gemini 是關(guān)鍵。
低成本微調(diào)的 LoRA
在微調(diào)部分，Colossal-AI 支持使用低秩矩陣微調(diào)（LoRA）方法。
LoRA 方法認(rèn)為大語言模型是過參數(shù)化的，其在微調(diào)中的參數(shù)改變量是一個低秩的矩陣，可以將其分解為兩個更小的的矩陣的乘積，即
在微調(diào)時，固定大模型參數(shù)，只調(diào)整低秩矩陣參數(shù)，從而顯著減小訓(xùn)練參數(shù)量。在微調(diào)之后，進(jìn)行推理部署之前，只需要將參數(shù)加回原有矩陣即可，即
，不增加模型的推理延遲。
減少內(nèi)存冗余的 ZeRO+Gemini
Colossal-AI 支持使用無冗余優(yōu)化器 (ZeRO) 來優(yōu)化內(nèi)存使用，這種方法可以有效減少內(nèi)存冗余，并且相比傳統(tǒng)的數(shù)據(jù)并行策略，不會犧牲計算粒度和通信效率，同時可以大幅提高內(nèi)存使用效率。
為了進(jìn)一步提升 ZeRO 的性能，Colossal-AI 引入了自動 Chunk 機制。
通過將運算順序上連續(xù)的一組參數(shù)存入同一個 Chunk 中（Chunk 是一段連續(xù)的內(nèi)存空間），可以確保每個 Chunk 的大小相同，從而提高內(nèi)存使用效率。
使用 Chunk 方式組織內(nèi)存可以保證 PCI-e 和 GPU-GPU 之間的網(wǎng)絡(luò)帶寬得到有效利用，減小通信次數(shù)，同時避免潛在的內(nèi)存碎片。
此外，Colossal-AI 的異構(gòu)內(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)存墻的限制，進(jìn)一步擴展了可訓(xùn)練模型規(guī)模。
One More Thing
盡管此次開源包含了復(fù)現(xiàn) ChatGPT 的完整算法流程和必要軟件系統(tǒng)，但想要走到實際應(yīng)用落地，還至少需要數(shù)據(jù)、算力等方面的支持。
參考開源大模型 BLOOM、開源 AI 畫畫工具 Stable Diffusion 的經(jīng)驗，這背后都需要包括個人開發(fā)者、算力、數(shù)據(jù)模型等可能合作方的支持共建 ——
此前，超過 1000 個科學(xué)家聯(lián)合發(fā)起、耗時一年多煉出了號稱和 GPT-3 一樣強大的語言模型 BLOOM。還有 AI 畫畫趨勢的頭號明星 Stable Diffusion，也是由 Stability AI、EleutherAI 和 LAION 多方聯(lián)合完成的。
復(fù)現(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)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。