首个开源 ChatGPT 低成本再现流程来了！预训练、奖励模型训练、强化学习训练，一次性打通。
明敏发自凹非寺，量子位 | 公众号QbitAI
最小demo训练流程仅需1.62GB显卡记忆体，随便一张消费级显卡就能满足。单卡模型容量最多提升10.3倍，相比原生PyTorch，单机训练速度最高可提升7.73倍，单卡推理速度提升1.42倍，仅需一行代码。
微调任务可最多提升单卡微调模型容量3.7倍，同时保持高速运行，同样仅需一行代码。ChatGPT红是红，再现也真的难，毕竟ChatGPT不开源，市面至今没有开源预训练权重、完全开源的低成本训练流程，且千亿等级大模型训练本身就是难题。
但ChatGPT军备赛已愈演愈烈，为了抓住趋势，如Google等都在打造竞品。快速再现ChatGPT是应趋势所需，开源加速方案Colossal-AI正是为此而来，且提供开源完整再现流程同时，也降低成本。
降低显卡记忆体消耗是关键
ChatGPT效果好，主要是训练过程引入人类回馈强化学习（RLHF），但也直接导致ChatGPT再现训练难度飙升。流程分为三阶段：
1. 监督微调：从Prompt库采样，收集人工回答，利用这些数据微调预训练大语言模型。
2. 奖励模型：从Prompt库采样，使用大语言模型产生多个回答，人工排序回答後，训练奖励模型（RM）模拟人类价值判断。
3. 基於第一阶段的监督微调模型和第二阶段奖励模型，利用强化学习演算法训练大语言模型。
▲ RLHF三个阶段。
对ChatGPT训练而言，第三阶段是核心部分。OpenAI采用强化学习近端策略最佳化演算法（PPO），引入奖励讯号，使语言模型产生内容更符合人类评判标准。但引入强化学习，代表有更多模型调用。如使用基於Actor-Critic（AC）结构的PPO演算法，需训练时进行Actor、Critic两个模型前向推理和反向传播，以及监督微调模型、奖励模型多次前向推理。
ChatGPT基础的InstructGPT论文，Actor和监督微调模型都使用1,750亿参数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能用原本一半硬体资源启动1,750亿参数模型训练，从64卡降低到32卡。如果续用64卡，则将训练时间压缩到更短，节省训练成本、加速产品更新。
为了让更多开发者体验再现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倍，还可继续扩大规模并行。
为了尽可能降低训练成本和上手门槛，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为系统策略即可快速启用。
使用下列指令即可快速启动单卡、单机多卡、1750亿版训练，并测试各种性能指标（包括最大记忆体占用、吞吐率和TFLOPS等）：
背後原理如何？
核心方案还是Colossal-AI。毕竟诞生起就是大模型应用专用，可基於PyTorch高效快速部署AI大模型训练和推理，是这领域的明星，GitHub Star超过8千，并成功入选SC、AAAI、PPoPP、CVPR等国际AI与HPC顶级会议的官方教程。…