HAI-Chat:高效且轻量的一键式 RLHF 训练工具
AI魔法学院
2024-03-25
分享海报

对话式人工智能工具的流行引发了人们对人工智能安全的担忧。构建有益、真实和无害行为的 LLMs,目前首选方案是基于人类反馈的强化学习(RLHF)。其中 ChatGPT 采用 RLHF 最先取得了显著的效果。

众所周知,RLHF 三阶段分别是:有监督微调(SFT),奖励模型(RM)与近端策略优化(PPO)。前两个步骤是基于特定的数据训练准备模型,属于人类反馈(HF)的部分;而 PPO 则是 OpenAI 提出的强化学习算法(RL),基于前两步构建的准备模型再训练提升主干网络。

幻方x深度求索在萤火智算集群上进行大量的 RLHF 训练实践,研发优化了一套轻量的 PPO 训练工具,名叫 HAI-Chat。它适配萤火智算集群的特性,更兼容HAI-LLM 大模型训练框架训练出来的模型,相比开源 RLHF 方案提速12倍。本文将为大家详细介绍。

RLHF 概述

回顾之前文章LLaMA-2 技术详解(一):数据打标》中所展示的 LLaMA-2 训练流程:

IMG_256

人反馈打标构建的数据信息通过 SFTRM  PPO 逐步融入到主干模型中。

步骤一:有监督微调

有监督微调(Supervised Fine-Tuning,简称 SFT),又被称为行为克隆(Behavioral Cloning,简称 BC),是 RLHF 训练过程中的一个重要步骤。在这个步骤中,预训练好的模型使用专家行为数据来进行微调训练,其旨在让模型的输出尽可能地复制或模仿专家的行为。这里不尝试优化某种奖励函数,因此它可能无法处理那些专家数据中没有覆盖到的情况,因为它完全依赖于专家的行为数据。

步骤二:奖励模型

奖励模型(Reward Model,简称 RM)是 RLHF 训练过程的第二阶段,它的目标是训练一个模型来适应人类的偏好。在这个阶段,首先从提示库中进行采样,并使用 LLM 生成多个响应。然后人工对这些响应进行排名,根据这些排名训练一个奖励模型。奖励模型的目标是学习人类对于不同响应的偏好,并将这些偏好编码到模型中。如下图所示:

这样,奖励模型可以用来为后续模型生成的新响应打分,从而在后续的训练中引导模型生成更符合人类偏好的内容。这种方式不仅能帮助模型处理训练数据中未覆盖的情况,也能减少模型生成不确定或模棱两可的回答,从而打破 SFT 行为克隆的影响。

在萤火智算集群中可以执行如下命令训练 RM

·        

hfai python train_reward.py --config hai_chat/configs/deepseek_rm_13b.py -- -n 8 -p 30 

步骤三:近端策略优化

近端策略优化(Proximal Policy Optimization,简称 PPO)是一种强化学习算法,它通过引入奖励信号来调整模型的行为,使模型生成的内容更符合人类的偏好。具体来说,PPO 通过最大化预期奖励来调整模型的策略,使模型在选择行为时更倾向于选择可以得到更高奖励的行为。如下图所示:

IMG_258

在这个阶段中,我们首先使用 SFT 微调的模型和奖励模型来生成一个初始的策略。然后,我们使用 PPO 算法来调整这个策略,使模型在生成内容时更考虑人类的偏好。通过这个阶段的训练,模型不仅可以理解人类的语言,还可以理解人类的偏好,并生成更符合人类偏好的内容。

在萤火智算集群中可以执行如下命令进行 PPO 训练:

·        

hfai python train_rlhf.py --config hai_chat/configs/deepseek_rlhf.py -- -n 2 -p 30

下面将为大家介绍 HAI-Chat PPO 的实现。

近端策略优化(PPO

具体来说,近端策略优化(PPO)可以分为三个阶段:

1.Rollout and Evaluation:在这个阶段,我们从 Prompt 库里抽样,使用语言模型生成文本,然后使用奖励模型(RM)给出奖励得分。这个得分反映了生成的文本的质量,比如它是否符合人类的偏好,是否符合任务的要求等。

2.Make Experience:在这个阶段,我们收集了一系列的经验,即模型的行为和对应的奖励。分为三步:旧策略采样(Old Policy Sampling),KL 散度惩罚(KL Penalty)与生成优势估计(Generalized Advantage Estimation,简称 GAE)。通过这三步,保证经过强化学习后的模型不会过于偏离原始预训练模型,提高整体的训练稳定性。

3.Optimization:这个阶段是 PPO 算法的关键阶段,包含五个步骤:新策略采样(New Policy Sampling),Critic LossActor LossEntropy Loss 和策略度量(PolicyKL)。在这个阶段,我们使用收集到的经验来更新模型的参数。具体来说,训练优化产生的新策略不断更新 Critic  Actor,加入熵和策略度量保证策略的多样性和训练稳定。

利用 HAI-Chat 自定义 PPO 的训练

HAI-Chat 提供了简单易用的接口,用户只需设置自定义的 tokenizer RLHF 训练配置(cfg),即可构建自己的 PPO 策略进行训练。示例代码如下:

1 from hai_chat.tokenizers.llama_tokenizer import LineBBPETokenizerForRLHF
2 from hai_chat.algorithms.ppo import PPO
3
4 tokenizer = LineBBPETokenizerForRLHF.get_tokenizer()
5
6 # 根据 cfg 构建 actor critic 模型与 dataloader
7 trainer = PPO(    
8      tokenizer=tokenizer,    
9      cfg=rlhf_cfg
10 )
11
12 # 加载模型
13 trainer.load_ckpt()
14
15 # 开始训练
16withtrainer.cm.start_save_thread(save_interval=rlhf_cfg.save_every_steps):    
17     for epoch in rlhf_cfg.epochs:        
18     trainer.set_epoch(epoch)        
19     trainer.train()

RLHF 训练配置(rlhf_cfg)示例如下:

这里 train() 函数里的执行逻辑是:

训练效果

对比目前最新的的 RLHF 开源训练框架,Safe-RLHFHAI-Chat 在相同配置下,相比 Safe-RLHF 速度提升了12倍。具体测试结果如下:

IMG_259

上图左侧展示了奖励模型每 Epoch 的训练与评估耗时,上图右侧展示了 PPO Epoch 的训练和 Rollout 耗时。可以看到,相比于 Safe-RLHFHAI-Chat 可以加速奖励模型 30%~50%,而对于 PPO 训练的加速则更加明显,可以达到10倍以上的加速。对于一个 7B 大小的模型在双机16 A100 上执行,Safe-RLHF 完整训练完成需要耗时 36 小时,而 HAI-Chat 只需 3 小时即可完成。HAI-Chat 可以大幅提高模型调优的速度,降低成本。

总结

本文简要介绍了 RLHF 的技术原理和 HAI-LLM 框架中 PPO 算法的实现。在大语言模型发展与商业化的道路中,RLHF 作为提升 LLMs 有益性和安全性的关键技术之一,将面对大量场景的挑战。我们将持续关注相关技术的改进与发展。

出自:https://mp.weixin.qq.com/s/v93EkKl9nDWoPR9-cRtdLA

© THE END

转载请联系本网站获得授权

投稿或版权问题请加微信:skillupvip