首个开源的 ChatGPT 低成本复现流程来了!
预训练、奖励模型训练、强化学习训练,一次性打通。
最小 demo 训练流程仅需 1.62GB 显存,随便一张消费级显卡都能满足了。单卡模型容量最多提升 10.3 倍。
相比原生 PyTorch,单机训练速度最高可提升 7.73 倍,单卡推理速度提升 1.42 倍,仅需一行代码即可调用。
对于微调任务,可最多提升单卡的微调模型容量 3.7 倍,同时保持高速运行,同样仅需一行代码。
要知道,ChatGPT 火是真的火,复现也是真的难。
毕竟 ChatGPT 是不开源的,市面上至今没有开源预训练权重、完全开源的低成本训练流程,而且千亿级别大模型的训练本身就是个难题。
但 ChatGPT 军备赛已经愈演愈烈,为了抓住趋势,如谷歌等都在打造对标竞品。快速复现 ChatGPT 是应趋势所需。
开源加速方案 Colossal-AI 正是为此而来。
并且在提供开源完整复现流程的同时,把成本降了下来!
降显存开销是关键
ChatGPT 的效果好,主要是由于在训练过程中引入了人类反馈强化学习,但这也直接导致 ChatGPT 的复现训练难度飙升。
其训练流程主要分为三个阶段:
1、监督微调:从 Prompt 库中采样,收集其人工回答,利用这些数据来微调预训练大语言模型;
2、奖励模型:从 Prompt 库中采样,使用大语言模型生成多个回答,人工对这些回答进行排序后,训练奖励模型,来拟合人类的价值判断。
3、基于第一阶段的监督微调模型和第二阶段的奖励模型,利用强化学习算法对大语言模型进一步训练。
RLHF 的三个阶段
对于 ChatGPT 训练而言,第三阶段是核心部分。
OpenAI 采用了强化学习中近端策略优化算法,借此引入奖励信号,使得语言模型生成内容更加符合人类评判标准。
但强化学习的引入,也意味着更多模型调用。
例如,使用基于 Actor-Critic结构的 PPO 算法,需要在训练时进行 Actor、Critic 两个模型的前向推理和反向传播,以及监督微调模型、奖励模型的多次前向推理。
在 ChatGPT 基础的 InstructGPT 论文中,Actor 和监督微调模型都使用了 1750 亿参数的 GPT-3 系列模型,Critic 和奖励模型则使用了 60 亿参数的 GPT-3 系列模型。
如此大规模的模型参数,意味着想要启动原始 ChatGPT 训练流程,需要数千 GB 的显存开销,单张 GPU 显然无法容纳,常见的数据并行技术也不能搞定。
即便引入张量并行、流水并行对参数进行划分,也需要至少 64 张 80GB 的 A100 作为硬件基础。而且流水并行本身并不适合 AIGC 的生成式任务,bubble 和调度复杂会导致效率受限。
单张消费级显卡都能体验
Colossal-AI 基于 ZeRO,Gemini, LoRA, Chunk-based 内存管理等方法,提出了一系列单卡、单机多卡、大规模并行解决方案。
对于基于 GPT-3 系列模型的 ChatGPT,Colossal-AI 能用原本一半的硬件资源启动 1750 亿参数模型训练,从 64 卡降低到 32 卡。
如果继续用 64 卡,则将训练时间压缩到更短,节省训练成本、加速产品迭代。
而为了能让更大范围的开发者体验复现 ChatGPT,除了 1750 亿参数版本外,Colossal-AI 还提供单卡、单机 4/8 卡的类 ChatGPT 版本,以降低硬件限制。
要知道,在单机多卡服务器上,即便把显卡规格提升为 A100 80GB,由于 ChatGPT 的复杂性和内存碎片,PyTorch 最大也只能启动基于 GPT-L这样的小模型 ChatGPT。
用 PyTorch 原生的 DistributedDataParallel 进行多卡并行扩展至 4 卡或 8 卡,性能提升有限。
Colossal-AI 最高可提升单机训练速度 7.73 倍,单卡推理速度 1.42 倍,还可继续扩大规模并行。
为了尽可能降低训练成本和上手门槛,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 开源方案,都能实现一行代码快速上手。
先看模型使用方面。
尽管 ChatGPT 背后的大语言模型 GPT-3.5 不开源,但如 GPT、OPT、BLOOM 等主流开源模型可作为替代。
Colossal-AI 为 Hugging Face 社区的这些模型,提供了开箱即用的 ChatGPT 复现代码,可覆盖三个阶段的训练。
以 GPT 为例,添加一行代码指定使用 Colossal-AI 作为系统策略即可快速使用。
fromchatgpt.nnimportGPTActor,GPTCritic,RewardModelfromchatgpt.trainerimportPPOTrainerfromchatgpt.trainer.strategiesimportColossalAIStrategywithstrategy.model_init_context:actor=GPTActor.cudacritic=GPTCritic.cudainitial_model=deepcopy(actor).cudareward_model=RewardModel(deepcopy(critic.model)).cudatrainer=PPOTrainer(strategy,actor,critic,reward_model,initial_model,...)trainer.fit(prompts)
使用下列命令,即可快速启动单卡、单机多卡、1750 亿版本训练,并测试各种性能指标:
#使用单机单卡训练GPT2-S,使用最小的batchsize,Colossal-AIGeminiCPU策略torchrun--standalone--nproc_pero_node1benchmark_gpt_dummy.py--models--strategycolossalai_gemini_cpu--experience_batch_size1--train_batch_size1#使用单机4卡训练GPT2-XL,使用Colossal-AIZero2策略torchrun--standalone--nproc_per_node4benchmark_gpt_dummy.py--modelxl--strategycolossalai_zero2#使用4机32卡训练GPT-3,使用Colossal-AIGeminiCPU策略torchrun--nnodes4--nproc_per_node8--rdzv_id=$JOB_ID--rdzv_backend=c10d--rdzv_endpoint=$HOST_NODE_ADDRbenchmark_gpt_dummy.py--model175b--strategycolossalai_gemini_cpu--experience_batch_
背后原理如何?
核心方案还是 Colossal-AI。
它从诞生起就面向大模型应用,可基于 PyTorch 高效快速部署 AI 大模型训练和推理,是这一领域的明星项目了,GitHub Star 超八千颗,并成功入选 SC、AAAI、PPoPP、CVPR 等国际 AI 与 HPC 顶级会议的官方教程。
目前,Colossal-AI 已成功帮助一家世界五百强企业,开发具备在线搜索引擎能力增强的类 ChatGPT 聊天机器人模型。
此前,它们还为 Stable Diffusion、OPT、AlphaFold 等前沿模型,提供了多样高效的大规模多维并行分布式解决方案。
主创人员为加州伯克利大学杰出教授 James Demmel 和新加坡国立大学校长青年教授尤洋。
具体到细节原理上,LoRA、ZeRO+Gemini 是关键。
低成本微调的 LoRA
在微调部分,Colossal-AI 支持使用低秩矩阵微调方法。
LoRA 方法认为大语言模型是过参数化的,其在微调中的参数改变量是一个低秩的矩阵,可以将其分解为两个更小的的矩阵的乘积,即
在微调时,固定大模型参数,只调整低秩矩阵参数,从而显著减小训练参数量。在微调之后,进行推理部署之前,只需要将参数加回原有矩阵即可,即
,不增加模型的推理延迟。
减少内存冗余的 ZeRO+Gemini
Colossal-AI 支持使用无冗余优化器 来优化内存使用,这种方法可以有效减少内存冗余,并且相比传统的数据并行策略,不会牺牲计算粒度和通信效率,同时可以大幅提高内存使用效率。
为了进一步提升 ZeRO 的性能,Colossal-AI 引入了自动 Chunk 机制。
通过将运算顺序上连续的一组参数存入同一个 Chunk 中,可以确保每个 Chunk 的大小相同,从而提高内存使用效率。
使用 Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU 之间的网络带宽得到有效利用,减小通信次数,同时避免潜在的内存碎片。
此外,Colossal-AI 的异构内存空间管理器 Gemini 支持将优化器状态从 GPU 卸载到 CPU ,以节省 GPU 内存占用。
可以同时利用 GPU 内存、CPU 内存来突破单 GPU 内存墙的限制,进一步扩展了可训练模型规模。
One More Thing
尽管此次开源包含了复现 ChatGPT 的完整算法流程和必要软件系统,但想要走到实际应用落地,还至少需要数据、算力等方面的支持。
参考开源大模型 BLOOM、开源 AI 画画工具 Stable Diffusion 的经验,这背后都需要包括个人开发者、算力、数据模型等可能合作方的支持共建 ——
此前,超过 1000 个科学家联合发起、耗时一年多炼出了号称和 GPT-3 一样强大的语言模型 BLOOM。还有 AI 画画趋势的头号明星 Stable Diffusion,也是由 Stability AI、EleutherAI 和 LAION 多方联合完成的。
复现 ChatGPT 也是如此,Colossal-AI 正在发起这一开发活动。
开源地址:
参考链接:
免责声明:该文章系本站转载,旨在为读者提供更多信息资讯。所涉内容不构成投资、消费建议,仅供读者参考。