复制ChatGPT的开源训练过程:Colossal-AI


ChatGPT是当今人工智能领域最热门的话题。 ChatGPT展示了非凡的能力,因此人们对复制它有很高的兴趣。

由于OpenAI一直没有发布ChatGPT的代码,如何有效复制ChatGPT就成了摆在大家面前的一个巨大问题,急需一个开源的ChatGPT等价物。

Colossal-AI刚刚开源了一个复制ChatGPT训练过程的解决方案,允许以极低的成本复制ChatGPT训练过程,并减少硬件限制。仅需 1.6GB GPU 内存即可使用,训练速度提高 7.73 倍!

亮点包括:

  • 一个开源完整的基于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 和许多其他预训练的大型模型到您的训练过程中。
  • 收敛性验证正在进行中,诚邀所有用户共同建设社区。
  • GitHub

ChatGPT技术分析
ChatGPT 取得惊人成绩的一个重要特点是在训练过程中引入了人类反馈强化学习(RLHF),以更好地捕捉人类的偏好。

ChatGPT 的训练过程分为三个主要阶段。

  1. 从 Prompt 库中抽样,收集其人类反应,并使用这些数据微调预训练的大型语言模型。
  2. 从 Prompt 库中采样,使用大型语言模型生成多个响应,手动对这些响应进行排序,并训练奖励模型 (RM) 以适应人类偏好。
  3. 基于第一阶段的监督微调模型和第二阶段的奖励模型,使用强化学习算法进一步训练大型语言模型。

在RLHF训练的核心部分stage 3,OpenAI采用强化​​学习中的Proximal Policy Optimization(PPO)算法引入奖励信号,使语言模型生成更符合人类偏好的内容。

ChatGPT 模型的复杂性源于强化学习的引入,会带来更多的模型调用。例如,使用基于Actor-Critic(AC)结构的PPO算法,我们需要在训练过程中对Actor和Critical模型进行前向推理和反向传播,以及有监督微调模型和奖励模型的多重前向推理。关于InstructGPT的论文,它是ChatGPT的基础,Actor和supervised fine-tuning模型都使用1750亿参数的GPT-3系列模型,而critical和reward模型使用60亿参数的GPT-3系列模型.

如此庞大的模型参数,要启动原有的ChatGPT训练过程, 需要上千GB的GPU显存,这显然远远超出了单个GPU的能力,普通的数据并行技术也不够用。然而,即使在分区参数中引入张量并行和流水线并行,仍然至少 需要64个80GB A100 GPU作为硬件基础。更糟糕的是,流水线不适合AIGC的生成任务,因为它的泡沫导致其复杂性和效率和调度。Stage 3涉及四种模型复杂的强化学习和训练过程,进一步给ChatGPT的代码复现带来困难和挑战。

使用 Colossal-AI 低成本复制 ChatGPT 训练过程
Colossal-AI开源复制了ChatGPT训练的基本流程,包括stage 1预训练,stage 2奖励模型训练,stage 3强化学习训练,这是流程中最复杂的阶段。

此外,Colossal-AI通过使用ZeRO、Gemini、LoRA、AutoChunk内存管理等大大降低了ChatGPT训练的GPU内存开销,只需要一半的硬件资源就可以启动1750亿参数模型训练(从64张卡到32张卡) ,显着降低了 ChatGPT 样式应用程序的成本。使用上述相同的硬件资源,Colossal-AI 能够以更短的时间进行训练,从而节省训练成本并加速产品迭代。

为了让更多开发者跑通ChatGPT训练流程,除了原有的1750亿参数版本,Colossal-AI还提供了高效的单GPU、独立4/8-GPUs ChatGPT-like版本,以减少硬件限制。

在单个多 GPUs 服务器上,即使使用最高端的 A100 80GB GPU,由于 ChatGPT 的复杂性和内存碎片,PyTorch 也只能基于 GPT-L(774M)等小型模型启动 ChatGPT。因此,使用 PyTorch 的 DistributedDataParallel (DDP) 将多 GPU 并行扩展到 4 或 8 个 GPU 会导致性能提升有限。

Colossal-AI 不仅在单 GPU 上的训练和推理加速方面具有显着优势,而且还可以随着并行性的扩大而进一步提升,单服务器训练速度最高可达 7.73 倍,单 GPU 推理速度最高可达 1.42 倍,并且 能够继续扩展到大规模并行,显着降低 ChatGPT 复制的成本。

此外,Colossal-AI 一直在努力降低基于预训练大型模型的微调任务的成本。例如,在 OPT 模型上关于 ChatGPT 的微调任务,Colossal-AI 能够在单 GPU 上将微调模型的容量提高到PyTorch 的3.7 倍 ,同时工作速度足够高。

一行代码启动
Colossal-AI 提供开箱即用的 ChatGPT 训练代码。用户会发现使用 Colossal-AI 以类似 ChatGPT 的方法训练来自 Hugging Face 社区的 GPT、OPT 和 BLOOM 等主流预训练模型很容易。例如,在 GPT 的情况下,只需要一行代码 来指定使用 Colossal-AI 作为系统策略来 kickstart。

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)

使用以下命令,可以快速启动单GPU规模、单机多GPU规模、原始1750亿参数规模版本的训练,并对各项性能指标(包括最大GPU内存占用、吞吐量和 TFLOPS):

# Training GPT2-S using a single card, a minimum batch size, Colossal-AI Gemini CPU strategy
torchrun --standalone --nproc_pero_node 1 benchmark_gpt_dummy.py --model s --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1
# Training GPT2-XL with a 4-GPU machine, Colossal-AI Zero2 strategy
torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2
# Training GPT-3 with 4 8-GPU servers, Colossal-AI Gemini CPU strategy
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_size 1 --train_batch_size 1

微信联系号:Colossal-AI

详细点击标题