明敏 發自 凹非寺

量子位 | 公眾號 QbitAI

這邊ChatGPT、GPT-4等AI大模型和應用打得火熱；

另一邊“平替”開源復現方案也加緊更新迭代。

這不，“首個開源ChatGPT低成本復現流程”就來了波大更新！

現在，僅需不到百億參數，利用RLHF簡單微調，模型即可掌握中、英雙語能力，達到與ChatGPT和GPT-3.5相當的效果。

中文對話效果如下：

這就是ColossalChat。

由Colossal-AI推出。一個月前，Colossal-AI乘著ChatGPT熱潮火速開源了低成本復現流程。

而新升級的ColossalChat，以Meta最新開源的LLaMA為基礎預訓練模型，能力更強、門檻還更低了：

• Demo：可直接在線體驗模型效果，無需注冊或waitinglist
• 訓練代碼：開源完整 RLHF 訓練代碼，已開源至含7B、13B兩種模型
• 數據集：開源104K中、英雙語數據集
• 推理部署：4bit量化推理70億參數模型僅需4GB顯存
• 模型權重：僅需單臺服務器少量算力即可快速復現
• 更大規模模型、數據集、其他優化等將保持高速迭代添加

要知道，模型開源、數據集、訓練應用成本、核心數據安全性等，是AI大模型浪潮下最被關注的一些問題。ColossalChat針對于此，一次性給出了應對方法。

由此，想要快速跟進ChatGPT這輪技術浪潮，門檻又低了一些。

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

包含完整 RLHF 流程

具體來看，ColossalChat的能力包括知識問答、中英文對話、內容創作、編程等。

和ChatGPT一樣，ColossalChat知道NBA和喬丹是誰：

會寫郵件，格式非常規范：

編程能力也不錯，能搞定算法：

據了解，為了能更逼近ChatGPT、GPT-4的驚艷效果，ColossalChat使用LLaMA作為預訓練模型，并包含完整RLHF流程。

之所以這樣做，主要是因為現有開源方案都可被視為只得到了人類反饋強化學習（RLHF）中第一步的監督微調模型，沒有進行后續的對齊和微調工作。

△RLHF的三個階段

比如Meta 開源了LLaMA模型，其參數量從70億到650億不等，號稱130 億參數即可勝過1750億的GPT-3模型在大多數基準測試的表現。

但是由于沒有被指令微調（instruct tuning），因此實際生成效果不夠理想。

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

而ChatGPT、GPT-4效果好的一大關鍵就是將RLHF引入訓練流程，才能讓生成內容更符合人類價值觀。

因此，基于LLaMA模型，包含完整RLHF流程的類Chat模型復現方案 ColossalChat，可以稱得上是目前最接近ChatGPT原始技術路線的實用開源項目。

總結來看，和Alpaca相比較ColossalChat具備4個方面的優勢。

第一、ColossalChat開源了第一個完整的RLHF pipeline，斯坦福Alpaca沒有做RLHF，也就是沒有做Stage2和3。

第二、ColossalChat采用了更多的指令數據，質量更好、范圍更大，并使用強化學習做了alignment使回答更接近人類。

第三、ColossalChat訓練流程集成了Colossal-AI的諸多系統優化，同等數據集和模型大小的訓練速度可以比AIpaca快3倍左右，能讓科研人員和中小企業獨立訓練部署自己的會話系統。

第四、ColossalChat采集了更多數據集：訓練的英文一共 24M tokens，中文大約 30M tokens，總共約 54M tokens。其中ColossalChat自己收集的數據集英文 6M，中文 18M tokens。

訓練數據集開源

數據集方面，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) 可提升3倍以上。

系統基礎設施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)，還可以進一步降低顯存要求，使用更低成本的硬件推理更大的模型。

開放協作

不過目前由于算力和數據集有限，在部分場景下的實際性能還有提升空間。

比如還是會被弱智吧問題難住：

而在這輪技術浪潮中，除了科技巨頭們，PyTorch、Hugging Face和OpenAI等開源社區與初創企業也起到了關鍵作用。

借鑒這些成功經驗，Colossal-AI也歡迎各方參與共建，并給出了多種參與方式：

• 在GitHub發布issue或提交pull request (PR)
• 加入Colossal-AI用戶微信或Slack群交流
• 發送正式合作提案到郵箱youy@comp.nus.edu.sg

如果你對這項工作感興趣，可以趕快和他們取得聯系~

開源地址：

https://github.com/hpcaitech/ColossalAI

參考鏈接：

https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b

— 完 —

量子位 QbitAI · 頭條號簽約