开源并行训练系统 ColossalAI 表示，已低成本复现了一个 ChatGPT 训练的基本流程，包括 stage 1 预训练、stage 2 的奖励模型的训练，以及最为复杂的 stage 3 强化学习训练。具体亮点包括：
一个开源完整的基于 PyTorch 的 ChatGPT 等效实现流程，涵盖所有 3 个阶段，可以帮助你构建基于预训练模型的 ChatGPT 式服务。
提供了一个迷你演示训练过程供用户试玩，它只需要 1.62GB 的 GPU 显存，并且可能在单个消费级 GPU 上实现，单 GPU 模型容量最多提升 10.3 倍。
与原始 PyTorch 相比，单机训练过程最高可提升 7.73 倍，单 GPU 推理速度提升 1.42 倍，仅需一行代码即可调用。
在微调任务上，同样仅需一行代码，就可以在保持足够高的运行速度的情况下，最多提升单 GPU 的微调模型容量 3.7 倍。
提供多个版本的单 GPU 规模、单节点多 GPU 规模和原始 1750 亿参数规模。还支持从 Hugging Face 导入 OPT、GPT-3、BLOOM 和许多其他预训练的大型模型到你的训练过程中。
ColossalAI 是一个具有高效并行化技术的综合大规模模型训练系统；旨在无缝整合不同的并行化技术范式，包括数据并行、管道并行、多张量并行和序列并行。其声称已通过 ZeRO、Gemini、Chunk-based 内存管理等技术，极大地降低 ChatGPT 训练的显存开销；仅需一半硬件资源即可启动 1750 亿参数模型训练（从 64 卡到 32 卡），显著降低应用成本。若使用上述相同硬件资源，Colossal-AI 则能以更短时间进行训练，节省训练成本，加速产品迭代。为了让更多开发者体验复现 ChatGPT 模型，除 1750 亿参数版本外，Colossal-AI 还提供高效的单 GPU、单机 4/8 GPU 的类 ChatGPT 版本，以降低硬件限制。博客内容还指出，在单机多 GPU 服务器上，即便使用最高端的 A100 80GB 显卡，由于 ChatGPT 的复杂性和内存碎片，PyTorch 最大仅能启动基于 GPT-L（774M）这样的小模型的 ChatGPT。用 PyTorch 原生的 DistributedDataParallel (DDP) 进行多卡并行扩展至 4 卡或 8 卡，性能提升有限。Colossal-AI 不仅在单 GPU 速度上训练和推理优势明显，随着并行规模扩大还可进一步提升，最高可提升单机训练速度 7.73 倍，单 GPU 推理速度 1.42 倍；并且能够继续扩展至大规模并行，显著降低 ChatGPT 复现成本。为了最大限度地降低培训成本和易用性，Colossal-AI 提供了可以在单个 GPU 上试用的 ChatGPT 培训流程。与在 14999 美元的 A100 80GB 上最多只能启动 7.8 亿个参数模型的 PyTorch 相比，Colossal-AI 将单个 GPU 的容量提升了 10.3 倍，达到 80 亿个参数。对于基于 1.2 亿参数的小模型的 ChatGPT 训练，至少需要 1.62GB 的 GPU 内存，任意单个消费级 GPU 都可以满足。此外，Colossal-AI 还在致力于降低基于预训练大型模型的微调任务的成本。以 ChatGPT 可选的开源基础模型 OPT 为例，Colossal-AI 能够在单 GPU 上将微调模型的容量提高到 PyTorch 的 3.7 倍，同时保持高速运行。Colossal-AI 为 Hugging Face 社区的 GPT、OPT 和 BLOOM 等主流预训练模型，提供了开箱即用的 ChatGPT 复现代码。以 GPT 为例，仅需一行代码，指定使用 Colossal-AI 作为系统策略即可快速使用。
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)
更多详情可查看博客内容：https://www.hpc-ai.tech/blog/colossal-ai-chatgpt