明敏 发自 凹非寺
量子位 | 公众号 QbitAI
这边ChatGPT、GPT-4等AI大模型和应用打得火热;
另一边“平替”开源复现方案也加紧更新迭代。
这不,“首个开源ChatGPT低成本复现流程”就来了波大更新!
现在,仅需 不到百亿参数,利用RLHF简单微调,模型即可掌握 中、英双语能力,达到与ChatGPT和GPT-3.5相当的效果。
中文对话效果如下:
这就是 ColossalChat。
由Colossal-AI推出。一个月前, Colossal-AI乘着ChatGPT热潮火速开源了低成本复现流程。
而新升级的 ColossalChat,以Meta最新开源的LLaMA为基础预训练模型,能力更强、门槛还更低了:
要知道,模型开源、数据集、训练应用成本、核心数据安全性等,是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),仅需一行代码即可使用。
ifargs.quant == '4bit':
model = load_quant(args.pretrained, args.gptq_checkpoint, 4, args.gptq_group_size)
如果采用高效的异步卸载技术(offload),还可以进一步降低显存要求,使用更低成本的硬件推理更大的模型。
开放协作
不过目前由于算力和数据集有限,在部分场景下的实际性能还有提升空间。
比如还是会被弱智吧问题难住:
而在这轮技术浪潮中,除了科技巨头们,PyTorch、Hugging Face和OpenAI等开源社区与初创企业也起到了关键作用。
借鉴这些成功经验,Colossal-AI也欢迎各方参与共建,并给出了多种参与方式:
如果你对这项工作感兴趣,可以赶快和他们取得联系~
开源地址:
https://github.com/hpcaitech/ColossalAI
参考链接:
https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b
— 完—
点这里 👇关注我,记得标星哦~
一键三连「分享」、「点赞」和「在看」
科技前沿进展日日相见 ~