简介
4月12日,微软宣布开源了 DeepSpeed Chat,帮助用户加速训练类似于 ChatGPT 的模型。
DeepSpeed Chat 能够简化 ChatGPT 类型模型的训练过程、强化推理体验。其中的 DeepSpeed-RLHF 系统能够在推理和训练两种模式之间进行切换,使复杂的 RLHF 训练速率更快、易于大规模推广。DeepSpeed-HE 比现有系统快 15 倍以上,且成本更低。在微软云 Azure 上,该系统只需 9 小时即可训练一个 OPT-13B 模型,只需 18 小时即可训练一个 OPT-30B 模型。
DeepSpeed-Chat 具有以下三大核心功能:
- 简化 ChatGPT 类型模型的训练和强化推理体验:只需一个脚本即可实现多个训练步骤,包括使用 Huggingface 预训练的模型、使用 DeepSpeed-RLHF 系统运行 InstructGPT 训练的所有三个步骤、甚至生成你自己的类 ChatGPT 模型。此外,还提供了一个易于使用的推理 API,用于用户在模型训练后测试对话式交互。
- DeepSpeed-RLHF 模块:DeepSpeed-RLHF 复刻了 InstructGPT 论文中的训练模式,并确保包括 a) 监督微调(SFT),b) 奖励模型微调和 c) 基于人类反馈的强化学习(RLHF)在内的三个步骤与其一一对应。此外,还提供了数据抽象和混合功能,以支持用户使用多个不同来源的数据源进行训练。
- DeepSpeed-RLHF 系统:其将 DeepSpeed 的训练(training engine)和推理能力(inference engine) 整合到一个统一的混合引擎(DeepSpeed Hybrid Engine or DeepSpeed-HE)中用于 RLHF 训练。DeepSpeed-HE 能够在 RLHF 中无缝地在推理和训练模式之间切换,使其能够利用来自 DeepSpeed-Inference 的各种优化,如张量并行计算和高性能 CUDA 算子进行语言生成,同时对训练部分还能从 ZeRO- 和 LoRA-based 内存优化策略中受益。DeepSpeed-HE 还能够自动在 RLHF 的不同阶段进行智能的内存管理和数据缓存。
DeepSpeed-Chat 具有以下三大特性:
-
效率和可负担性:在效率方面,DeepSpeed-HE比现有系统快15倍以上,使RLHF培训既快速又实惠。例如,DeepSpeed-HE可以在Azure Cloud上以低于13美元和9美元的价格在30小时内训练OPT-18B,在300小时内训练OPT-600B。
-
出色的可扩展性:DeepSpeed-HE支持数千亿参数的模型,可以在多节点多GPU系统上实现出色的可扩展性。因此,即使是 13B 模型也可以在 1.25 小时内完成训练,而使用 DeepSpeed-HE 可以在一天内训练一个庞大的 175B 模型。
上面两个表中的数字是训练的第 3 步,基于 DeepSpeed-RLHF 策划数据集和训练配方上实际测量的训练吞吐量,该训练配方在总共 135M 个代币上训练一个时期。我们总共有 67.5M 个查询令牌(131.9k 个查询,序列长度为 256)和 67.5M 个生成的令牌(131.9k 个答案,序列长度为 256),每步的最大全局批大小为 0.5M 个令牌(1024 个查询-答案对)。
-
RLHF 训练民主化:只需一个 GPU,DeepSpeed-HE 就支持超过 13 亿个参数的训练模型,使数据科学家无需访问多 GPU 系统即可创建不仅可以创建玩具 RLHF 模型,还可以创建可用于现实场景的大型强大模型。
上手体验
示例1:提供从DeepSpeed-Chat训练的1.3B最终模型的问答环节
示例 2:通过提供从 DeepSpeed-Chat 训练的模型进行多轮对话
一个脚本完成 RLHF 训练的所有三个阶段,并生成您的第一个 ChatGPT 模型!
pip install deepspeed>=0.9.0
git clone https://github.com/microsoft/DeepSpeedExamples.git
cd DeepSpeedExamples/applications/DeepSpeed-Chat/
pip install -r requirements.txt
python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node
在大约半天的时间里,13B 模型将完全训练并准备好检查点。下表显示了三个步骤中每个步骤的训练时间明细:
不同的型号尺寸和配置
例如,如果您想在 GPU 集群上为您的研究或业务训练更大、更高质量的模型,您可以简单地使用相同的脚本来构建所需的模型大小,例如 66B 和 GPU 计数,例如 64 个 GPU
python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node
在 9 小时内,您就可以准备好 66 亿个参数 ChatGPT 模型,以便在您最喜欢的前端 GUI 中提供服务:
如果你只有大约1-2个小时的咖啡或午休时间,你也可以尝试用DeepSpeed-Chat训练一个小型/玩具模型。例如,我们为具有单个数据集的 1.3B 模型准备了一个训练示例,以便在您的消费级 GPU 上测试。
python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu
具体来说,针对不同规模的模型和硬件配置,DeepSpeed-RLHF系统所需的时间和成本如下:
使用DeepSpeed-Chat的RLHF API自定义您自己的RLHF训练管道
DeepSpeed-Chat允许用户使用如下所示的灵活API构建自己的RLHF训练管道,用户可以使用这些API来重建自己的RLHF训练策略。这为创建用于研究探索的各种RLHF算法提供了通用接口和后端。
engine = DeepSpeedRLHFEngine(
actor_model_name_or_path=args.actor_model_name_or_path,
critic_model_name_or_path=args.critic_model_name_or_path,
tokenizer=tokenizer,
num_total_iters=num_total_iters,
args=args)
trainer = DeepSpeedPPOTrainer(engine=engine, args=args)
for prompt_batch in prompt_train_dataloader:
out = trainer.generate_experience(prompt_batch)
actor_loss, critic_loss = trainer.train_rlhf(out)
参考:https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-chat/README.md
参考:https://www.deepspeed.ai/
参考:https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/README.md#-deepspeed-chats-roadmap-
成熟的 RLHF 培训管道
为了提供无缝的培训体验,我们遵循 InstructGPT,并在 DeepSpeed-Chat 中包含成熟的端到端训练管道,如图 1 所示。
我们的管道包括三个主要步骤:
- 第 1 步:监督微调 (SFT),其中人类对各种查询的响应经过精心挑选,以微调预训练的语言模型。
- 第 2 步:奖励模型微调,其中使用具有人工提供的同一查询的多个答案排名的数据集训练单独的(通常小于 SFT)模型 (RW)。
- 第 3 步:RLHF 训练,其中使用近端策略优化 (PPO) 算法通过来自 RW 模型的奖励反馈进一步微调 SFT 模型。
我们在步骤 3 中提供了两个附加功能来帮助提高模型质量:
- 指数移动平均线 (EMA) 集合,可以选择基于 EMA 的检查点进行最终评估。
- 混合训练,将预训练目标(即下一个单词预测)与 PPO 目标混合在一起,以防止在 SQuAD2.0 等公共基准测试上出现回归性能。
两个训练功能(EMA 和混合训练)经常被其他最近的工作省略,因为它们可以是可选的。然而,根据InstructGPT,EMA检查点通常比传统的最终训练模型提供更好的响应质量,混合训练可以帮助模型保持训练前的基准求解能力。因此,我们为用户提供它们,以充分获得 InstructGPT 中所述的培训体验,并争取更高的模型质量。
除了与InstructGPT论文高度一致外,我们还提供了方便的功能,以支持研究人员和从业者使用多种数据资源训练自己的RLHF模型:
- 数据抽象和混合功能:DeepSpeed-Chat能够使用多个数据集训练模型,以获得更好的模型质量。它配备了(1)抽象数据集层,以统一不同数据集的格式;(2)数据拆分/混合功能,以便正确混合多个数据集,然后在3个训练阶段中拆分。
DeepSpeed 混合引擎 – 统一的基础设施,为RLHF培训提供动力和优化
指令引导的 RLHF 流水线的步骤 1 和步骤 2 类似于大型模型的常规微调,它们由基于 ZeRO 的优化和 DeepSpeed 训练中并行策略的灵活组合提供支持,以实现规模和速度。另一方面,就性能影响而言,管道的步骤 3 是最复杂的部分。每次迭代都需要有效处理两个阶段:a)令牌/体验生成的推理阶段,为训练生成输入,b)训练阶段更新参与者和奖励模型的权重,以及它们之间的交互和调度。它引入了两个主要成本:(1)内存成本,因为在整个阶段3中需要提供SFT和RW模型的多个副本;(2)主要发电阶段,如果加速不当,将大大减慢整个阶段3。此外,我们在第 3 阶段添加的两个重要功能(包括指数移动平均线 (EMA) 收集和混合训练)将产生额外的内存和训练成本。
为了应对这些挑战,我们将DeepSpeed Training和Inference的完整系统功能组合到一个统一的基础设施中,我们称之为混合引擎。它利用原始的DeepSpeed引擎进行快速训练模式,同时毫不费力地将DeepSpeed推理引擎应用于生成/评估模式,为第3阶段的RLHF训练提供了明显更快的训练系统。如图 2 所示,DeepSpeed 训练和推理引擎之间的过渡是无缝的:通过为执行组件模型启用典型的评估和训练模式,在运行推理和训练管道时,DeepSpeed 会选择其不同的优化来更快地运行模型并提高整体系统吞吐量。