本文经AI新媒体量子位授权转载（公众号ID：QbitAI），转载请联系出处用于转载。首款开源ChatGPT低成本复现流程来了！一次完成预训练、奖励模型训练和强化学习训练。最小的demo训练过程只需要1.62GB显存，任何消费级显卡都可以满足。单卡机型容量提升高达10.3倍。与原生PyTorch相比，单机训练速度最高可提升7.73倍，单卡推理速度可提升1.42倍。只需一行代码即可调用。对于微调任务，单卡微调模型容量最高可提升3.7倍，同时保持高速运行，仅需一行代码。要知道，ChatGPT真的很流行，要复制它真的很难。毕竟，ChatGPT不是开源的。到目前为止，市面上还没有开源的预训练权重和完全开源的低成本训练过程，千亿级大模型的训练本身就是一个问题。然而，ChatGPT军备竞赛愈演愈烈。为了抓住趋势，谷歌等正在打造竞品。快速重现ChatGPT是顺应潮流所必需的。开源加速解决方案Colossal-AI就是为此而来。并且在提供完整的开源复现过程的同时，降低了成本！开源地址：https://github.com/hpcaitech/ColossalAI降低内存成本是关键ChatGPT之所以能运作良好，主要得益于在训练过程中引入了人类反馈强化学习（RLHF），但这也直接导致了再次出现ChatGPTTraining难度飙升。其训练过程主要分为三个阶段：1.监督微调：从Prompt库中抽样，收集其人工答案，利用这些数据对预训练的大语言模型进行微调；2.奖励模型：从Prompt库中抽样，使用大语言模型生成多个答案，人工对这些答案进行排序后，训练奖励模型（RM）来拟合人类的价值判断。3.基于第一阶段的监督微调模型和第二阶段的奖励模型，使用强化学习算法进一步训练大语言模型。△RLHF的三个阶段对于ChatGPT的训练，第三阶段是核心部分。OpenAI在强化学习中采用近端策略优化算法（PPO），从而引入奖励信号，使语言模型生成的内容更符合人类的判断标准。但是强化学习的引入也意味着更多的模型调用。例如，使用基于Actor-Critic（AC）结构的PPO算法，需要在训练过程中对Actor和Critic模型进行前向推理和反向传播，以及监督微调模型和奖励模型的多次前向推理。在基于ChatGPT的InstructGPT论文中，Actor和supervisedfine-tuning模型均采用1750亿参数的GPT-3系列模型，Critic和reward模型采用60亿参数的GPT-3系列模型。如此大规模的模型参数意味着需要数千GB的显存开销才能启动原始的ChatGPT训练过程，这显然不是单个GPU所能容纳的，普通的数据并行技术也无法应对。即使引入张量并行和流水线并行来划分参数，也至少需要64块80GB的A100作为硬件基础。而且pipeline并行本身并不适合AIGC的生成任务，bubble和调度的复杂性会导致效率有限。单卡消费级显卡体验Colossal-AI基于ZeRO、Gemini、LoRA、基于Chunk的内存管理等方式，提出一系列单卡、单机多卡、大规模并行解决方案。对于基于GPT-3系列模型的ChatGPT，Colossal-AI可以使用原来一半的硬件资源启动1750亿参数模型训练，从64张卡片减少到32张卡片。如果继续使用64张卡，训练时间会缩短，节省训练成本，加速产品迭代。为了让更广泛的开发者体验到ChatGPT的再现，除了1750亿参数版本，Colossal-AI还提供了单卡、单机4/8卡式ChatGPT，以减少硬件限制。要知道在单机多卡服务器上，即使显卡规格升级到A10080GB，由于ChatGPT的复杂性和内存碎片，PyTorch也只能启动基于GPT-L的小模型ChatGPT(774M)。使用PyTorch原生的DistributedDataParallel（DDP）进行多卡并行扩展到4卡或8卡，性能提升有限。Colossal-AI最高可将单机训练速度提升7.73倍，单卡推理速度提升1.42倍，并可继续扩大并行规模。为了尽可能降低训练成本和进入门槛，Colossal-AI还提供了ChatGPT训练流程，可以在单GPU上试用。与10万元左右的A10080GB最多只能启动一个模型7.8亿参数的PyTorch相比，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为HuggingFace社区的这些模型提供了开箱即用的ChatGPT复现代码，可以涵盖训练的三个阶段。以GPT为例，添加一行代码指定使用Colossal-AI作为系统策略即可快速使用。fromchatgpt.nnimportGPTActor,GPTCritic,RewardModelfromchatgpt.trainerimportPPOTrainerfromchatgpt.trainer.strategiesimportColossalAIStrategystrategy=ColossalAIStrategy(stage=3,placement_policy='cuda')withstrategy.model_init_context()(PT.Actor(actor)=)复制代码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等）：#使用单机单卡训练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-AIZero2strategytorchrun--standalone--nproc_per_node4benchmark_gpt_dummy.py--modelxl--strategycolossalai_zero2#使用4台机器32张卡训练GPT-3，使用Colossal-AI双子座niCPUstrategytorchrun--nnodes4--nproc_per_node8\--rdzv_id=$JOB_ID--rdzv_backend=c10d--rdzv_endpoint=$HOST_NODE_ADDR\benchmark_gpt_dummy.py--model175b--strategycolossalai_gemini_cpu--experience_batch_背后的原理是什么它？核心方案是Colossal-AI，从诞生之日起就面向大规模模型应用。可以高效快速部署基于PyTorch的AI大规模模型训练和推理。AAAI、PPoPP、CVPR等国际顶级AI和HPC会议官方教程。目前，Colossal-AI已成功帮助一家世界500强公司开发了一种具有增强的在线搜索引擎功能的类ChatGPT聊天机器人模型。此前，他们还为StableDiffusion、OPT、AlphaFold等前沿模型提供了多种高效的大规模多维并行分布式解决方案。主要创作者为美国加州大学伯克利分校特聘教授JamesDemmel和新加坡国立大学青年教授YouYang。△Colossal-AI与当今同期主要开源项目的开源数据对比在细节和原理上，LoRA、ZeRO+Gemini是重点。用于低成本微调的LoRA在微调部分，Colossal-AI支持使用低秩矩阵微调(LoRA)方法。LoRA方法认为大语言模型是过参数化的，它在fine-tuning中的参数变化是一个低秩矩阵，可以分解为两个更小的矩阵的乘积，即fine-tuning时，大模型是固定参数，只调整低秩矩阵参数，从而显着减少训练参数量。微调后，推理部署前，只需将参数加回原矩阵即可，即不增加模型的推理延迟。△LoRA示意图，只需要训练A和B就可以减少内存冗余ZeRO+GeminiColossal-AI支持使用非冗余优化器（ZeRO）来优化内存使用，这种方法可以有效减少内存冗余，并且与传统的相比数据并行策略不牺牲计算粒度和通信效率，并且可以大大提高内存使用效率。为了进一步提升ZeRO的性能，Colossal-AI引入了自动分块机制。通过将一组按运算顺序连续的参数存储到同一个Chunk中（一个Chunk是一个连续的内存空间），可以保证每个Chunk的大小相同，从而提高内存使用效率。使用Chunk来组织内存可以保证PCI-e和GPU-GPU之间网络带宽的有效利用，减少通信次数，避免潜在的内存碎片。△Chunk机制此外，Colossal-AI的异构内存空间管理器Gemini支持将优化器状态从GPU卸载到CPU，以节省GPU内存使用。GPU内存和CPU内存（由CPUDRAM或NVMeSSD内存组成）可以同时使用，突破单个GPU内存墙的限制，进一步扩大可训练模型的规模。△通过ZeRO+GeminiOneMoreThing增强硬件的模型能力虽然本次开源包含了完整的算法流程和复现ChatGPT所必需的软件系统，但要实现实际应用，至少需要数据和算力的支持。参考开源大模型BLOOM和开源AI绘图工具StableDiffusion的经验，这些都需要个体开发者、算力、数据模型等可能合作伙伴的支持和共建。-此前，1000多位科学家共同发起并耗时一年多时间完善语言模型BLOOM，号称与GPT-3一样强大。还有AI绘画潮流的头号明星StableDiffusion，也是由StabilityAI、EleutherAI和LAION共同完成的。由Colossal-AI发起开发的ChatGPT的复制也是如此。如果你对这项工作感兴趣或有好的想法，可以通过以下方式与他们取得联系：在GitHub上发布issue或提交PR，加入Colossal-AI用户微信或Slack群交流，点击“阅读原文”原文”填写合作提案，发送合作提案至邮箱contact@hpcaitech.com传送门：开源地址：https://github.com/hpcaitech/ColossalAI
