大模型RLHF的trick

来自:包包算法笔记

进NLP群—>加入NLP交流群

之前在文章大模型面试八股中提到一个问题,大模型训练中RM分数越来越高,那训出来LLM的效果一定好吗?

这么肯定的判断肯定是有坑的,值得怀疑。

如果你动手跑几次ppo的过程就发现了,大模型的强化学习非常难以训练,难以训练不仅仅指的是费卡,还是指的非常容易训崩。

第一,费卡。假设你训llama 7b,SFT 和 RM 都用7B的模型,那么显存耗费 = 2*7B(TRIAN MODE) + *7B(EVAL MODE), 分别对应 policy model / critic model,还有ref model/reward model

本来你能用几张40GB A100的卡+deepspeed 做7b的全参数微调,强化学习就得升级到80GB的A100了,勉勉强强能跑到7B。想跑更大的就得充钱了。

第二,容易崩。LLM训着训着就不听你话了,要么变成停不下来的复读机,输出到后面没有逻辑直到maxlen,要么变成哑巴,直接一个eosid躺平。

RLHF中的问题其实在RL游戏训练里面很常见了,如果环境和参数设置不好的话,agent很容走极端,在 一头撞死or循环鬼畜之间反复横跳。

原始的ppo就是很难训,对SFT基模型和RM的训练数据以及采样prompt的数据要求很高,参数设置要求也很高。

自从openai带了一波RLHF的节奏后,大家都觉得强化学习在对齐方面的无敌功力,但自己跑似乎又不是那么回事,这玩意也太有讲究了吧。

更多的魔鬼在细节了,openai像拿了一个比赛的冠军,告诉你了成功的solution,结果没告诉你各个步骤的重要性和关键设置,更没有告诉你失败和无效的经验。

这里要夸一下复旦,复旦一篇非常实在的技术报告提出了7个trick,不仅仅告诉你关键的trick设置,也告诉你失败的经验,如果不怎么样模型就会一头撞死。

参考https://github.com/OpenLMLab/MOSS-RLHF

在讲trick之前,首先复旦-MOSS也对LLM的训练过程加了更多监测,其实这些都是实验中非常重要的监控过程指标,能很清楚的发现你模型是否出现异常。

然后这个图很好,非常清楚地讲述了trick是如何作用在RLHF中的各个阶段的,另外配套的开源代码实现也非常清晰易懂,典型的面条代码没有什么封装,一码到底,易读性和魔改都很方便。

下面我们看看这7个trick,对应图中右侧画星号的部分。

7820340c6ca59e91195d9c85f60e49ad.png

1, token级别的KL散度惩罚

kl_penalty = (-self.kl_penalty_weight * (logprobs - ref_logprobs)).cpu()

这一步主要解决的问题是训练稳定性,防止步子太大扯着蛋,如果你输出的和参考模型差别过大的话就减分。

2,Reward Normalization and Clipping

3,Value Function Loss Clipping

Clipping类似梯度裁剪,也是止步子太大扯着蛋,对一些异常的loss和reward做了限制,Normalization为了对reward做标准化。

这部分的代码可以对应开源中的这些设置仔细查看,原理大同小异

self.use_reward_clip: bool = opt.use_reward_clip
self.use_reward_norm: bool = opt.use_reward_norm
self.use_advantage_norm: bool = opt.use_advantage_norm
self.use_advantage_clip: bool = opt.use_advantage_clip
self.use_critic_loss_clip: bool = opt.use_critic_loss_clip
self.use_policy_loss_clip: bool = opt.use_policy_loss_clip

4.Critic Model Initialization

用RM model初始化Critic可能不是一个必要的选择,作者做了一些实验证明这个问题,推荐使用critic model pre-training。代码里这部分还没有,还是使用rm初始化的,后续跟进一下这个问题。

5. Generalized Advantage Estimation

附录里C.3有GAE的调参实验。

8ef747bd9203d4c3d0693c6151270eb9.png

6.Clipped Surrogate Objective

这个也是一种正则化方法,防止步子太大扯着蛋,确保训练过程的中的稳定性,这个方法比一般policy gradient处理的更为高效。

d3189569e9a9e26eedb0d63f7643667a.png

7.Global Gradient Clipping

ddc8dec80de02c6c1a1779d6b522feb4.png

原理还是同上,所有的Clipping无非都是砍掉太大的步子。

另外作者还用了一个instruct gpt里面用到的方案,增加了训练过程使用 llm_pretrain_loss,参考代码

if self.use_entropy_loss:
    loss1 = pg_loss + self.vf_loss_weight * vf_loss + self.entropy_loss_weight * entro_loss
else:
    loss1 = pg_loss + self.vf_loss_weight * vf_loss
loss2 = self.ppo_pretrain_loss_weight * pretrain_loss
loss = loss1 + loss2

总结下,整体ppo-max的改进主要集中在训练过程的稳定性上,用的东西还是模型的老三样,训练过程裁剪,初始化,loss改进,主要集中在如何能让RLHF更好调,推荐参考作者的源码进行一些实验。

另外,作者在论文里留了一个彩蛋,技术报告的第二部分预告主要是讲Reward Model的成功和踩坑经验,目前还没有发布,静待作者更新。之前大家一直的争论点用什么scale的RM,说要用远远大于SFT model的RM model,这到底是不是一个关键的问题,是不是deberta 和 65B都行,期待作者第二个技术报告里给一个实验~


进NLP群—>加入NLP交流群

猜你喜欢

转载自blog.csdn.net/qq_27590277/article/details/131778090
今日推荐