开源方案复现ChatGPT流程！1.62GB显存即可体验，单机训练提速7.73倍

语言: CN / TW / HK
时间 2023-02-15 16:02:40
主题:

火爆全网的ChatGPT，仿佛开启了第四次工业革命，让微软、谷歌等全球科技巨头打得昏天黑地，引得各路玩家纷纷入局，抢占赛道。

然而由于OpenAI没有开源ChatGPT，如何有效复现ChatGPT已成为摆在大家面前的头号难题，急需可靠的开源共建方案。

Colossal-AI快速跟进，首个开源低成本复现ChatGPT完整流程。作为当下最火热的开源AI大模型解决方案，Colossal-AI已收获开源社区GitHub Star近万颗，此次开源亮点包括：

  • 开源完整基于PyTorch的ChatGPT复现流程，涵盖全部3个阶段，可实现从预训练模型到ChatGPT的蜕变；
  • 体验最小demo训练流程最低仅需1.62GB显存，任意单张消费级GPU即可满足，单卡模型容量最多提升10.3倍；
  • 相比原生PyTorch，最高可提升单机训练速度7.73倍，单卡推理速度1.42倍，一行代码即可使用；
  • 对于微调任务，可最多提升单卡的微调模型容量3.7倍，同时保持高速运行，仅需一行代码；
  • 提供单卡、单机多卡、1750亿参数等多个版本，支持从Hugging Face导入OPT，GPT-3，BLOOM等多种预训练大模型；
  • 收敛验证正在进行中，该项目也在吸引合作者共建生态。

开源地址：http://github.com/hpcaitech/ColossalAI

 

ChatGPT——AIGC引发的工业革命

如果问新年伊始，最火爆的科技热点是什么？非ChatGPT莫属。

它仿佛无所不能的六边形战士，可以聊天、写代码、修改 bug、做表格、发论文、写作业、做翻译、甚至代替Google搜索引擎等……

自发布以来，ChatGPT便已摧枯拉朽之势席卷各个行业，不仅5天时间便突破百万用户，月活用户突破1亿更是仅用时2个月，成为史上增速最快的消费级应用，远超如今其他知名应用，如Twitter 5 年、Meta（Facebook）4 年半，TikTok 9个月等，而手机普及到1亿用户则用了16年。


               1亿用户月活用户耗时
                                    1亿用户月活用户耗时

比尔·盖茨盛赞“ChatGPT的意义不亚于PC和互联网诞生”，而微软CEO萨蒂亚·纳德拉（Satya Nadella）更是直言“堪比工业革命，这辈子第一次见这么大的技术浪潮” 和“AI 正在重塑互联网”。作为向OpenAI投资上百亿美元的大金主，微软已火速将ChatGPT整合进自家的搜索引擎必应Bing和Edge浏览器，还计划加入Teams以及Office等办公套件全家桶，股价一夜市值飙涨超800亿美元。


               
                                      微软与谷歌发布会后股价对比

而隔壁需要担心被ChatGPT革命掉自家搜索引擎的谷歌，虽然拉响“红色警报”，紧急发布对标竞品Bard，却因Demo首秀翻车，股价市值瞬间蒸发1000亿美元。

一夜之间，全球的科技巨头们仿佛都回到了自己年轻时的样子，纷纷宣布要打造自己的ChatGPT。

但ChatGPT发布已有数月，市面上不仅没有预训练权重开源，连可靠的完整开源训练流程都仍是空白，更无法实现基于千亿大模型的ChatGPT全流程高效搭建和应用。临时上线，号称“对标ChatGPT”的一众新品们，因为闭源也难辨真伪。

为什么ChatGPT有如此魔力？复现它又有哪些难点？

 

ChatGPT技术分析

ChatGPT的惊人效果，重要特征是在训练过程引入人类反馈强化学习（RLHF），使得模型表现更符合人类价值观。

ChatGPT的训练流程主要分为三个阶段：

  1. 从Prompt库中采样，收集其人工回答，利用这些数据来微调预训练大语言模型。
  1. 从Prompt库中采样，使用大语言模型生成多个回答，人工对这些回答进行排序后，训练奖励模型（RM），来拟合人类的价值判断。
  1. 基于阶段1的监督微调模型和阶段2的奖励模型，利用强化学习算法对大语言模型进一步训练。

其中阶段3是RLHF训练的核心部分，OpenAI采用了强化学习中的近端策略优化算法（PPO），借此引入奖励信号，使得语言模型生成内容更加符合人类评判标准。


              
                                           RLHF的三个阶段

ChatGPT模型的复杂性在于强化学习的引入会带来更多模型的调用。例如，使用基于Actor-Critic（AC）结构的PPO算法，需要在训练时进行Actor、Critic两个模型的前向推理和反向传播，以及监督微调模型、奖励模型的多次前向推理。在ChatGPT基础的InstructGPT的论文中，Actor和监督微调模型都使用了1750亿参数的GPT-3系列模型，Critic和奖励模型则使用了60亿参数的GPT-3系列模型。

对于如此多的模型参数，想要启动原始ChatGPT训练流程，需要数千GB的显存开销，显然远超单张GPU的容纳能力，常见的数据并行技术也无能为力。但即使引入张量并行、流水并行对参数进行划分，也仍需至少64张80GB的A100作为硬件基础。并且，其中的流水并行由于bubble和调度复杂，效率受限，不适合AIGC的生成式任务。阶段3涉及4个模型的复杂强化学习训练流程，进一步给ChatGPT的代码复现带来了困难和挑战。

使用Colossal-AI低成本复现ChatGPT

Colossal-AI以开源方式复现了ChatGPT训练的基本流程，包括阶段1预训练，阶段2的奖励模型的训练，以及最为复杂的阶段3的强化学习训练等。

同时，Colossal-AI通过ZeRO，Gemini,  Chunk-based内存管理等技术，极大地降低ChatGPT训练的显存开销，仅需一半硬件资源即可启动1750亿参数模型训练（64卡->32卡），显著降低应用成本。若使用上述相同硬件资源，Colossal-AI则能以更短时间进行训练，节省训练成本，加速产品迭代。

为了让更多开发者体验复现ChatGPT模型，除1750亿参数版本外，Colossal-AI还提供高效的单卡、单机4/8卡的类ChatGPT版本，以降低硬件限制。


               

 

在单机多卡服务器上，即便使用最高端的A100 80GB显卡，由于ChatGPT的复杂性和内存碎片，PyTorch最大仅能启动基于GPT-L（774M）这样的小模型的ChatGPT。用PyTorch原生的DistributedDataParallel (DDP) 进行多卡并行扩展至4卡或8卡，性能提升有限。

Colossal-AI不仅在单卡速度上训练和推理优势明显，随着并行规模扩大还可进一步提升，最高可提升单机训练速度7.73倍，单卡推理速度1.42倍，还可继续扩展至大规模并行，显著降低ChatGPT复现成本。


                   

为了尽可能降低训练成本和上手门槛，Colossal-AI还提供了在单张GPU上即可尝试的ChatGPT训练流程。相比于PyTorch在约10万元的A100 80GB上，最大仅能启动7.8亿参数模型，Colossal-AI将单卡容量提升10.3倍至80亿参数。对于基于1.2亿参数小模型的ChatGPT训练，最低仅需1.62GB显存，任意单张消费级GPU即可满足。


             

此外，Colossal-AI也致力于降低基于预训练大模型的微调任务成本。以ChatGPT可选的开源基础模型OPT为例，相比PyTorch，Colossal-AI可将提升单卡微调模型容量3.7倍（原始计算量显著增大），同时保持高速运行。

一行代码快速上手

Colossal-AI为Hugging Face社区的GPT，OPT和BLOOM等主流预训练模型，提供了开箱即用的ChatGPT复现代码。以GPT为例，仅需一行代码，指定使用Colossal-AI作为系统策略即可快速使用。

Python
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)

使用下列命令，即可快速启动单卡、单机多卡、1750亿版本训练，并测试各种性能指标（包括最大显存占用、吞吐率和TFLOPS等）：

Python
# 使用单机单卡训练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卡训练GPT2-XL，使用Colossal-AI Zero2策略
torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2
# 使用4机32卡训练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_size 1 --train_batch_size 1

背后优化

  • 核心系统Colossal-AI

复现ChatGPT的背后，依赖面向大模型时代的通用深度学习系统Colossal-AI，可基于PyTorch高效快速部署AI大模型训练和推理，降低AI大模型应用成本。

自开源以来，Colossal-AI已经多次在GitHub热榜位列世界第一，获得GitHub Star超八千颗，并成功入选SC、AAAI、PPoPP、CVPR等国际AI与HPC顶级会议的官方教程。除上述优化外，Colossal-AI还针对AI大模型趋势，提供最多样和高效的大规模多维并行分布式解决方案，此前已在Stable Diffusion、OPT、AlphaFold等前沿模型上展现卓越优势。


           
                   Colossal-AI与当今主要开源项目同期开源数据对比

Colossal-AI由加州伯克利大学杰出教授James Demmel和新加坡国立大学校长青年教授尤洋领导。相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用，广受好评。Colossal-AI已成功帮助某世界500强企业，开发具备在线搜索引擎能力增强的类ChatGPT聊天机器人模型。

  • 低成本微调的LoRA

Colossal-AI支持使用低秩矩阵微调（LoRA）方法进行高效微调。LoRA方法认为大语言模型是过参数化的，其在微调中的参数改变量是一个低秩的矩阵，可以将其分解为两个更小的的矩阵的乘积，即。在微调时，固定大模型参数，只调整低秩矩阵参数，从而显著减小训练参数量。在微调之后，进行推理部署之前，只需要将参数加回原有矩阵即可，即, 不增加模型的推理延迟。
                 
                             LoRA示意图，仅需训练A、B

  • 减少内存冗余的ZeRO + Gemini

Colossal-AI 支持使用无冗余优化器 (ZeRO) 来优化内存使用，这种方法可以有效减少内存冗余，并且相比传统的数据并行策略，不会牺牲计算粒度和通信效率，同时可以大幅提高内存使用效率。为了进一步提升 ZeRO 的性能，Colossal-AI 引入了自动Chunk机制。通过将运算顺序上连续的一组参数存入同一个 Chunk中（Chunk 是一段连续的内存空间），可以确保每个 Chunk 的大小相同，从而提高内存使用效率。使用Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU之间的网络带宽得到有效利用，减小通信次数，同时避免潜在的内存碎片。


             
                                                     Chunk机制

此外，Colossal-AI的异构内存空间管理器Gemini支持将优化器状态从 GPU 卸载到 CPU ，以节省 GPU 内存占用。可以同时利用 GPU 内存、CPU 内存（由 CPU DRAM 或 NVMe SSD内存组成）来突破单GPU内存墙的限制，进一步扩展了可训练模型规模。


                 
                                   通过ZeRO + Gemini提升硬件的模型容量

开放协作

尽管此次开源包含了复现ChatGPT的完整算法流程和必要软件系统，但对于像ChatGPT这样的超大AI大模型，想要实际落地应用，还需要数据、算力至少2方面的努力。毕竟训练一个1750亿参数的GPT-3就需要数百万美元算力。因此，长期以来预训练大模型都由少数大型私营科技公司垄断。

好在开源社区已成功进行了新的尝试。例如，完全开放代码、数据集、权重的1760亿参数的BLOOM模型，共有来自全球60个国家、超过250个机构，以及超过1000名研究人员参与其中，其中包括以个人名义参加的Meta、谷歌等大厂员工。而前段时间大火的开源图文生成模型Stable Diffusion，也是由Stability AI、EleutherAI和LAION等组织共同完成的。

借鉴上述成功模式，该项目也在吸引更多的合作者：无论是个人开发者，还是算力、数据、模型等可能合作方，都有机会参与其中，大显身手，以复现ChatGPT为起点，拥抱大模型时代！

可通过以下方式联系或参与：

  • 在GitHub发布issue或提交PR
  • 加入Colossal-AI用户微信或Slack群交流
  • 点击阅读原文填写合作提案
  • 发送合作提案到邮箱[email protected]

 

开源地址：http://github.com/hpcaitech/ColossalAI

 

参考链接：http://www.hpc-ai.tech/blog/colossal-ai-chatgpt

 

（阅读原文链接）http://www.hpc-ai.tech/partners

 

「」