出品|开源中国

开源并行训练系统 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, RewardModelfrom chatgpt.trainer import PPOTrainerfrom chatgpt.trainer.strategies import ColossalAIStrategystrategy = 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)