火爆全網的 ChatGPT,完整流程方案,現已開源!

在過去的短短几個月,以 ChatGPT、GPT4 爲代表的 AI 應用和大模型火爆全球,被視爲開啓了新的科技工業革命和 AGI (通用人工智能)的新起點。

不僅科技巨頭間你追我趕,爭相推出新品,許多學術界、工業界的 AI 大佬也紛紛投入投身相關創業浪潮。

生成式 AI 正以 “天” 爲單位,快速迭代,持續狂飆!

然而,OpenAI 並未將其開源,它們背後的技術細節有哪些?如何快速跟進、追趕並參與到此輪技術浪潮中?如何降低 AI 大模型構建和應用的高昂成本?如何保護核心數據與知識產權不會因使用第三方大模型 API 外泄?

作爲當下最受歡迎的開源 AI 大模型解決方案,GitHub 知名開源項目 Colossal-AI 率先建立了包含監督數據集收集 -> 監督微調 -> 獎勵模型訓練 -> 強化學習微調的完整 RLHF 流程。

以 LLaMA 爲基礎預訓練模型,推出 ColossalChat,是目前最接近 ChatGPT 原始技術方案的實用開源項目!

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

包含以下內容

  1. Demo:可直接在線體驗模型效果,無需註冊或 waitinglist;

  2. 訓練代碼:開源完整 RLHF 訓練代碼,已開源至含 7B 和 13B 兩種模型;

  3. 數據集:開源 104K 中、英雙語數據集;

  4. 推理部署:4bit 量化推理 70 億參數模型僅需 4GB 顯存;

  5. 模型權重:僅需單臺服務器少量算力即可快速復現;

  6. 更大規模模型、數據集、其他優化等將保持高速迭代添加。

平價模型,強大能力

ColossalChat 僅需不到百億參數,在大語言模型的基礎上進行 RLHF 微調,即可掌握中、英雙語能力,達到與 ChatGPT 和 GPT-3.5 類似的效果。

例如常識問答:    

中文應答:

寫一封郵件:

寫個算法:

完整 ChatGPT 克隆方案

儘管 ChatGPT 和 GPT-4 等 GPT 系列模型非常強大,但是它們不太可能被完全開源。幸運的是,開源社區一直在不斷努力。

例如 Meta 開源了 LLaMA 模型,該模型的參數量從 70 億到 650 億不等,130 億參數即可勝過 1750 億的 GPT-3 模型在大多數基準測試的表現。但是由於沒有被指令微調(instruct tuning),因此實際生成效果不夠理想。

斯坦福的 Alpaca 通過調用 OpenAI API,以 self-instruct 方式生成訓練數據,使得僅有 70 億參數的輕量級模型以極低成本微調後,即可獲得媲美 GPT-3.5 這樣千億參數的超大規模語言模型的對話效果。

但是現有開源方案都可以被視爲只得到了人類反饋強化學習(RLHF)中第一步的監督微調模型,沒有進行後續的對齊和微調工作。同時 Alpaca 的訓練數據集過小,語料只有英文,也在一定程度上限制了模型的性能。

ChatGPT 和 GPT-4 的驚豔效果,還在於將 RLHF 引入訓練過程,使得生成內容更加符合人類價值觀。

RLHF 的三個階段

基於 LLaMA 模型,Colossal-AI 首個開源包含完整 RLHF 流程的類 Chat 模型復現方案 ColossalChat,是目前最接近 ChatGPT 原始技術路線的實用開源項目!

訓練數據集開源

ColossalChat 開源了包含約 10 萬條問答的中、英雙語數據集。該數據集收集並清洗了社交平臺上人們的真實提問場景作爲種子數據集,利用 self-instruct 技術擴充數據,花費約 900 美元進行標註。

對比其他 self-instruct 方法生成的數據集,該數據集的種子數據更加真實、豐富,生成的數據集涵蓋的話題更多。該數據可以同時用於微調和 RLHF 訓練。通過高質量的數據,ColossalChat 能進行更好地對話交互,同時支持中文。

ColossalChat 數據集收集流程

RLHF 算法復現

RLHF-Stage1 是 supervised-fintuning,即使用上文提到的數據集進行模型微調。

RLHF-Stage2 訓練了獎勵模型,它通過對於同一個 prompt 的不同輸出進行人工排序,得到對應分數,監督訓練獎勵模型。

RLHF-Stage3 使用了強化學習算法,是訓練流程中最複雜的一部分:

RLHF-Stage3 算法流程圖

在 PPO 部分,ColossalChat 分爲兩個階段進行:首先是 Make Experience 部分,利用 SFT 、Actor、RM、Critic 模型計算生成 Experience 存入 buffer 中;之後是參數更新部分,利用 Experience 計算策略損失和價值損失。

在 PTX 部分,ColossalChat 計算 Actor 輸出 response 和輸入語料的回答部分的交叉熵損失函數,用來在 PPO 梯度中加入預訓練梯度,以保持語言模型原有性能防止遺忘。最後將策略損失、價值損失和 PTX 損失加和進行反向傳播和參數更新。

快速上手

ColossalChat 開源了基於 LLaMA 模型,復現訓練 ChatGPT 三個階段的完整代碼。

第一階段,訓練 SFT 模型:

# Training with a 4-GPU servers
colossalai run --nproc_per_node=4 train_sft.py \
   --pretrain "/path/to/LLaMa-7B/" \
   --model 'llama' \
   --strategy colossalai_zero2 \
   --log_interval 10 \
   --save_path /path/to/Coati-7B \
   --dataset /path/to/data.json \
   --batch_size 4 \
   --accimulation_steps 8 \
   --lr 2e-5

第二階段,訓練獎勵模型:

# Training with a 4-GPU servers
colossalai run --nproc_per_node=4 train_reward_model.py \
   --pretrain "/path/to/LLaMa-7B/" \
   --model 'llama' \
   --strategy colossalai_zero2 \
   --dataset /path/to/datasets

第三階段,使用 RL 訓練:

# Training with a 8-GPU servers
colossalai run --nproc_per_node=8 train_prompts.py prompts.csv \
   --strategy colossalai_zero2 \
   --pretrain "/path/to/Coati-7B" \
   --model 'llama' \
   --pretrain_dataset /path/to/dataset

在獲得最終模型權重後,還可通過量化降低推理硬件成本,並啓動在線推理服務,僅需單張約 4GB 顯存的 GPU 即可完成 70 億參數模型推理服務部署。

python server.py /path/to/pretrained --quant 4bit --gptq_checkpoint /path/to/coati-7b-4bit-128g.pt --gptq_group_size 128

系統性能優化與開發加速

ColossalChat 能夠快速跟進 ChatGPT 完整 RLHF 流程復現,離不開 AI 大模型基礎設施 Colossal-AI 及相關優化技術的底座支持,相同條件下訓練速度相比 Alpaca 採用的 FSDP(Fully Sharded Data Parallel) 可提升兩倍以上

系統基礎設施 Colossal-AI

AI 大模型開發系統 Colossal-AI 爲該方案提供了基礎支持,它可基於 PyTorch 高效快速部署 AI 大模型訓練和推理,從而降低 AI 大模型應用的成本。

Colossal-AI 由加州伯克利大學傑出教授 James Demmel 和新加坡國立大學校長青年教授尤洋領導開發。

自開源以來,Colossal-AI 已經多次在 GitHub 熱榜位列世界第一,獲得 GitHub Star 約兩萬顆,併成功入選 SC、AAAI、PPoPP、CVPR、ISC 等國際 AI 與 HPC 頂級會議的官方教程。

減少內存冗餘的 ZeRO + Gemini

Colossal-AI 支持使用無冗餘優化器 (ZeRO) 提高內存使用效率,低成本容納更大模型,同時不影響計算粒度和通信效率。自動 Chunk 機制可以進一步提升 ZeRO 的性能,提高內存使用效率,減少通信次數並避免內存碎片。

異構內存空間管理器 Gemini 支持將優化器狀態從 GPU 顯存卸載到 CPU 內存或硬盤空間,以突破 GPU 顯存容量限制,擴展可訓練模型的規模,降低 AI 大模型應用成本。

使用 LoRA 低成本微調

Colossal-AI 支持使用低秩矩陣微調(LoRA)方法,對 AI 大模型進行低成本微調。LoRA 方法認爲大語言模型是過參數化的,而在微調時,參數改變量是一個低秩矩陣。

因此,可以將這個矩陣分解爲兩個更小的矩陣的乘積。在微調過程中,大模型的參數被固定,只有低秩矩陣參數被調整,從而顯著減小了訓練所需的參數量,並降低成本。

低成本量化推理

GPTQ 量化

爲降低推理部署成本,Colossal-AI 使用 GPTQ 4bit 量化推理。在 GPT/OPT/BLOOM 類模型上,它比傳統的 RTN(rount-to-nearest) 量化技術能夠獲得更好的 Perplexity 效果。

相比常見的 FP16 推理,它可將顯存消耗降低 75%,只損失極少量的吞吐速度與 Perplexity 性能。

以 ColossalChat-7B 爲例,在使用 4bit 量化推理時,70 億參數模型僅需大約 4GB 顯存即可完成短序列(生成長度爲 128 )推理,在普通消費級顯卡上即可完成(例如 RTX 3060 Laptop),僅需一行代碼即可使用。

if args.quant == '4bit':
   model = load_quant(args.pretrained, args.gptq_checkpoint, 4, args.gptq_group_size)

如果採用高效的異步卸載技術 (offload),還可以進一步降低顯存要求,使用更低成本的硬件推理更大的模型。

開放協作

儘管已經進一步引入 RLHF,但由於算力和數據集有限,在部分場景下的實際性能仍有提升空間。

幸運的是,不同以往 AI 大模型與前沿技術僅由少數科技巨頭壟斷,PyTorch、Hugging Face 和 OpenAI 等開源社區與初創企業在本輪浪潮中也起到了關鍵作用。

借鑑開源社區的成功經驗,Colossal-AI 歡迎各方參與共建,擁抱大模型時代!

可通過以下方式聯繫或參與:

  1. 在 GitHub 發佈 issue 或提交 pull request (PR)

  2. 加入 Colossal-AI 用戶微信或 Slack 羣交流

  3. 發送正式合作提案到郵箱 youy@comp.nus.edu.sg

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

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/SLiqxn2AbfVLtp1MAYEg_Q