首个开源的 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），但这也直接导致 ChatGPT 的复现训练难度飙升。

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

1、监督微调：从 Prompt 库中采样，收集其人工回答，利用这些数据来微调预训练大语言模型；

2、奖励模型：从 Prompt 库中采样，使用大语言模型生成多个回答，人工对这些回答进行排序后，训练奖励模型（RM），来拟合人类的价值判断。

3、基于第一阶段的监督微调模型和第二阶段的奖励模型，利用强化学习算法对大语言模型进一步训练。

△RLHF 的三个阶段

对于 ChatGPT 训练而言，第三阶段是核心部分。

OpenAI 采用了强化学习中近端策略优化算法（PPO），借此引入奖励信号，使得语言模型生成内容更加符合人类评判标准。

但强化学习的引入，也意味着更多模型调用。

例如，使用基于 Actor-Critic（AC）结构的 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（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 作为系统策略即可快速使用。

from chatgpt.nn import GPTActor, GPTCritic, RewardModel
from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import ColossalAIStrategy



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 等）：

# 使用单机单卡训练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_

背后原理如何？

核心方案还是 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）方法。

LoRA 方法认为大语言模型是过参数化的，其在微调中的参数改变量是一个低秩的矩阵，可以将其分解为两个更小的的矩阵的乘积，即

在微调时，固定大模型参数，只调整低秩矩阵参数，从而显著减小训练参数量。在微调之后，进行推理部署之前，只需要将参数加回原有矩阵即可，即

，不增加模型的推理延迟。

减少内存冗余的 ZeRO+Gemini

Colossal-AI 支持使用无冗余优化器 (ZeRO) 来优化内存使用，这种方法可以有效减少内存冗余，并且相比传统的数据并行策略，不会牺牲计算粒度和通信效率，同时可以大幅提高内存使用效率。

为了进一步提升 ZeRO 的性能，Colossal-AI 引入了自动 Chunk 机制。

通过将运算顺序上连续的一组参数存入同一个 Chunk 中（Chunk 是一段连续的内存空间），可以确保每个 Chunk 的大小相同，从而提高内存使用效率。

使用 Chunk 方式组织内存可以保证 PCI-e 和 GPU-GPU 之间的网络带宽得到有效利用，减小通信次数，同时避免潜在的内存碎片。

此外，Colossal-AI 的异构内存空间管理器 Gemini 支持将优化器状态从 GPU 卸载到 CPU ，以节省 GPU 内存占用。

可以同时利用 GPU 内存、CPU 内存（由 CPU DRAM 或 NVMe SSD 内存组成）来突破单 GPU 内存墙的限制，进一步扩展了可训练模型规模。

One More Thing

尽管此次开源包含了复现 ChatGPT 的完整算法流程和必要软件系统，但想要走到实际应用落地，还至少需要数据、算力等方面的支持。

参考开源大模型 BLOOM、开源 AI 画画工具 Stable Diffusion 的经验，这背后都需要包括个人开发者、算力、数据模型等可能合作方的支持共建 ——

此前，超过 1000 个科学家联合发起、耗时一年多炼出了号称和 GPT-3 一样强大的语言模型 BLOOM。还有 AI 画画趋势的头号明星 Stable Diffusion，也是由 Stability AI、EleutherAI 和 LAION 多方联合完成的。

复现 ChatGPT 也是如此，Colossal-AI 正在发起这一开发活动。

开源地址：

https://github.com/hpcaitech/ColossalAI

参考链接：

• https://www.hpc-ai.tech/blog/colossal-ai-chatgpt

本文来自微信公众号：量子位 （ID：QbitAI），作者：明敏