关于Instruct GPT复现的一些细节与想法
关于Instruct GPT复现的一些细节与想法
最近有机会参与到一个InstructGPT的复现工作中来，写一篇文章做个简单的介绍；
part 1 简单介绍RLHF，看过太多遍的朋友可以跳过；
part 2 简单介绍了我们实现的一些细节，如果感兴趣的话可以之后再讲讲具体的代码结构；
part 3 提出了一些我的不成熟的想法和问题，欢迎大家讨论~
1 简单介绍一下RLHF
ChatGPT具有比传统语言模型更出色的效果，这很大程度上归因于采用了人类反馈强化学习方法（Reinforcement Learning from Human Feedback, RLHF）的训练模式。该训练模式将人类评价引入了大语言模型的训练，以使得基于一般文本语料训练得到的大语言模型更符合人类价值观念。
ChatGPT的训练流程主要分为三个阶段。第一阶段，从Prompt数据库中采样，收集其人工回答，利用这些数据来微调预训练大语言模型。第二阶段， 从Prompt库中采样，使用大语言模型生成多个回答，人工对这些回答进行排序，利用排序信息，训练奖励模型（RM），来拟合人类对于生成语言的价值判断。第三阶段，基于第一阶段的监督微调模型（SFT）和第二阶段的奖励模型，利用强化学习算法（如PPO)进行大语言模型的进一步训练。
ChatGPT模型的复杂性在于它采用了强化学习方法，这极大地增加了模型的调用复杂度。例如我们使用PPO这种基于Actor-Critic（AC）结构的强化学习算法，我们就需要在训练时进行Actor、Critic两个模型的前向推理和反向传播，以及监督微调模型、奖励模型的前向推理。在InstructGPT的论文中，Actor和监督微调模型都使用了1750亿参数的GPT3模型，Critic和奖励模型都使用了60亿参数的GPT3模型。简单算一下的话，ChatGPT的训练过程最少需要（175+6）*（16+4）= 3620 GB的显存，需要最少6个8卡80GB A100的节点（还是没计算其他开销的情况下），还是比较夸张的。复杂的强化学习训练流程给ChatGPT的代码复现带来了很大的难度，巨大的模型参数量也给ChatGPT的训练和推理都带来了极大的挑战。
2 简单说说复现
由于ChatGPT论文并没有放出来，复现的基础也就只能是InstructGPT了。
2.1 Step2 - Reward Model的训练
奖励模型的训练也是RLHF流程重要的一个环节，这一部分实际上是一个监督训练过程。 首先我们从Prompt数据集中取出一些，然后输入初始模型（SFT），对于一个prompt产生多个response，具体怎么产生呢？可以是多个语言模型，也可以是generate之后采样，或者变一下随机种子什么的，都可以，只要确保这多个response是同一个prompt生成的就好。然后利用人工评价，来对于这一些同一个prompt的response排名，利用排名来计算出（prompt，response）二值对的分数。InstructGPT论文里使用的是pair-wise loss来将rank计算为具体的分数，当然也可以用经典的ELO算法来实现。得到分数，就作为sequence的lable，可以对于RM进行监督训练。
2.2 Step3-RL
RLHF最主要的第三阶段使用了PPO算法，主要流程是可以分两步走的，make experience和train。
在这个算法中，我们认为输入的prompt是state，输出的response是action，想要得到的策略就是怎么从prompt生成action能够得到最大的reward，也就是拟合人类的偏好。
我下面简单地讲解一下这两个部分的思路。
2.2.1 Make Experience
在Make Experience的流程中可以分两步，全都只涉及推理，不涉及反向传播。
一共涉及四个模型，Actor，Critic，Supervised Fine-tune Model（SFT）， Reward Model（RM）；其中Actor和SFT都应该是175B的模型，Critic和RM是6B的模型。其中Actor参数由SFT初始化，Critic参数由RM初始化，其模型结构分别对应相同。
首先是从Prompt数据库里采样出来，然后输入SFT，做generate，得到（prompt + response）= sequence。这一个过程中我们需要从prompt批量采样来进行generate，得到一个sequence buffer。
然后，我们将sequence buffer内容batched输入四个模型做inference，Actor输出action_logits，Critic输出value（标量）,RM输出r(x, y)（标量），SFT输出sft_logits；sft_logits和action_logits做kl散度，为了约束actor模型的更新step不要偏离原始模型太远；再由r(x, y)和kl_div计算得到reward；reward和value计算得到adv。我们将sequence、action_logits、value、adv、reward打包，放进replay buffer，make experience的流程就结束了。
# 获取sequence
sequence_buffer = []
for prompt in prompt_dataset:
sequence= actor.generate(prompt)
sequence_buffer.append(sequence)
# 获取exp
exp = []
for sequence in sequence_buffer:
action_logits = actor(sequence)
sft_action_logits = sft(sequence)
value = critic(sequence)
reward = reward_model(sequence) - kl_coef * KL_func(action_logits.logprob，sft_action_logits.logprob)
# adv怎么算我还有点疑惑
adv = reward - value
exp.append([action, sequence, action_logits, value, reward, adv])
2.2.2 Train
这一部分就是训练部分，会更新actor和critic的参数，也是吃显存的大头。
Actor的流程是取出sequence，然后inference生成新的logits，再和sequence对应的之前的logits计算ratio，和adv计算出pg_loss，也就是actor的loss，然后反向传播，优化器迭代。
Critic的流程是取出sequence，然后inference得到新的value，和old_value做clip_value，和reward计算value loss，然后反向传播，优化器迭代。
action_logits = actor(exp['sequence'])
ratio = (action_logits.logprob - exp['action_logits'].logprob).exp()
surr1 = ratio * exp['adv']
surr2 = retio.clip(1 - eps_clip, 1 + eps_clip) * exp['adv']
policy_loss = -torch.min(surr1, surr2)
value = critic(exp['sequence'])
value_clip = exp['value'] + (value - exp['value']).clamp(-v_clip, v_clip)
surr1 = (value_clip - exp['reward']) ** 2
surr2 = (value - exp['reward']) ** 2
value_loss = max(surr1, surr2)
2.2.3 分布式RL
现在make experience和train两个部分已经有了，我们该怎么组装这两个组件呢？
一个想法是顺序执行，先make experience，把replay buffer填满，然后再训练，直到replay buffer清空，再继续make experience。优点是占用显存相较于第二个想法会小一些。
另一个想法是分布式并行，这个分布式不是TP之类的分布式技术，而是稍微高一点维度的（其实就是分布式RL，类似于distributed - ppo的思路），可以使用多个experience maker生成数据，然后trainer消耗数据，控制experience产生和消耗的速度基本一致，就可以维护一个稳定的replay buffer，然后模型定期从trainer更新到experiecn maker。优点是效率应该会比较高，但是大规模部署175B模型的推理和训练，即使使用了ZeRO优化器、异构计算等技术，还是需要极其高的成本的。（之前做过175B大模型的并行量化推理，可以之后说说）
至此，基本简述了InstructGPT的复现思路和一些细节，更多的细节欢迎讨论。也欢迎关注Colossal-AI[4]的复现及训练方案。
ColossalAI/applications/ChatGPT at main · hpcaitech/ColossalAI (github.com)
3 一些问题的讨论
3.1 关于critic
其实复现过程中我一度误入歧途，错误地理解了这个公式：
InstructGPT论文[2]中优化函数的公式
可能是太久没有碰强化学习了，竟然混淆了objective funtion和loss function，真是罪过。这个错误导致我一开始以为论文说的ppo只是一个壳子，只是用了importance sampling来限制更新的步长；并没有使用actor-critic结构，只是用了一个actor而已。
当时还一通分析，说“value function（critic）存在的意义在于拟合环境的reward生成，来判定一个动作的好坏，那么是不是有了reward model，把答案都放在这里了，就没必要再去训练、拟合了value funciton了呢？“虽然后来我通过论文中一些其他的线索，推断出来肯定是存在critic的，论文的训练流程使用的就是最经典的PPO算法，但是这个疑问还是没有解开，希望有理解的朋友可以给我指点指点。
3.2 关于step
我们应该将一次generate过程认为是一个step，还是一次inference过程认为是一个step（如果不了解GPT的朋友可以认为，inference一次生成一个词，generate是多次inference，直到生成了结束符或者到达了max_len）？从InstructGPT论文的角度讲，它的设定应该是认为这是一个one-step的过程，每一个step生成一个完整的句子，而reward model也是为了评判每一个（prompt，response）对的好坏，而没有办法评判每一个词生成的好坏。我有一个问题：“那这样不就是浪费了强化学习解决顺序决策问题的能力？“
而且，生成式语言模型的generate过程，完全可以认为是一个多step的episode，有非常合适的state的转移的描述：从n个词长度的prompt到n+1个词长度的prompt；有一个终局reward：reward model的输出。那我们现在需要的就是更稠密的reward反馈，即每一步的reward，可以思考一下该如何实现。在强化学习的方法都帮助下，语言模型的生成可以考虑到更多的“未来信息”，也就是当前token生成对于之后token的影响，我个人认为是对于生成过程更有帮助的。
还有一种可能，或许我们可以把视角拉大，认为每一次完整的对话是一个episode，每一句的对话认为是一个step，这样就可以给出一个考虑了顺序决策的RL建模。但是问题在于，你给出reward的时候，是需要考虑前文的，这样又好像不太符合马尔可夫性，除非认为所有之前的对话都是state。有想法的朋友欢迎讨论。
0227更新：感谢@思悥的指正，这里我的理解存在了比较大的问题，正确的版本应该是：
每个token的生成（inference）认为是一个step，每一次完整句子的生成（generate）认为是一个episode；
每一个action是一个token的生成，action space 是词表空间，state在一个step之后更新为(state+new_token)；
每个step的reward是kl_div，每一个episode end的reward 是reward-model的输出。
3.3 一些小细节的问题
- adv的计算不用GAE怎么算呢，我们的实现是 reward - value，我觉得不一定对；
- pair-wise loss的计算，我理解中算期望时候已经做过一次平均了，那么前面再平均一次不会有问题吗？
3. Entropy loss需要算吗？根据openai之前文章的一些描述，我们的实现没有计算entropy loss。
4. 感谢直播时一个同学的提问，KL散度计算时，kl散度应该是一个非对称的，那么两个logits哪个应该在前，哪个在后呢？我个人理解，下图公式计算的是分布P相对于Q的相对熵，抽象一点来看，Q应该是已知的，P应该是我们需要衡量的，因此我觉得代码中使用D_kl(action_logits, old_action_logits) 是合理的。
4 总结
刚好之前做过一些强化学习相关的工作，现在又在做mlsys，所以有机会亲手实现InstructGPT的复现。
以上就是我的一些理解，还有一些想法和不少问题，欢迎朋友们友善讨论。
另外感谢HPC-AI（Colossal-AI）的支持，对于ChatGPT有更多想法的朋友们欢迎参与到Colossal-AI相关的开源工作中来 :-），有colossal-ai相关问题的同学欢迎评论区讨论，或者去github提issue反馈，谢谢大家！
Reference
[1]OpenAI博客：https://openai.com/blog/chatgpt/
[2]InstructGPT论文：Training language models to follow instructions with human feedback
[3]Huggingface博客：Illustrating Reinforcement Learning from Human Feedback (RLHF)
[4]Colossal-AI: GitHub - hpcaitech/ColossalAI: Making big AI models cheaper, easier, and scalable
[5]Learning to summarize from human feedback（OpenAI稍早的工作）：Learning to summarize from human feedback
[6]KL_div的wikipedia：Kullback–Leibler divergence - Wikipedia