本文为《类ChatGPT逐行代码解读》系列的第二篇，上一篇是：

本文模型的特点是都加了RLHF，对于这4个模型而言：TRL、ChatLLaMA、ColossalChat、DeepSpeed Chat

• 如果只关注两个 则可以更多关注下ColossalChat、DeepSpeed Chat，原因在于ColossalChat给的图示特别好，而DeepSpeed Chat的实现很清晰
• 如果有读者说 就只想看一个，则推荐DeepSpeed Chat(简称DSC)，特别是DSC会给你一个完整而通透的“PPO算法/RLHF”的代码实现全流程，好的资料可以让你事半功倍

总之，微软这个DeepSpeed Chat实现的不错，抠完它的关键代码后，你会发现和之前本博客内另一篇写的原理部分都一一对应起来了(如果你还没看过原理，建议先看此文：，只有懂原理才能更好的理解实现或实际实现，特别是该文的第三部分 )，而把论文、原理/算法、公式、代码一一对应，可以让你的理解有个质变

本文最早的标题是：从零实现带RLHF的类ChatGPT：从TRL/ChatLLaMA/ColossalChat到DeepSpeed Chat，后来因为要不断扩展DSC的内容，为避免本文越写越长，故最终分出了两篇文章

1. 本文侧重：TRL、ChatLLaMA、ColossalChat
2. 新文侧重：

通过《ChatGPT技术原理解析》一文，我们已经知道了ChatGPT的三阶段训练过程，其中，阶段三的本质其实就是通过PPO的方式去微调LM

GitHub上有个TRL(Transformer Reinforcement Learning，基于『Hugging Face开发的Transformer库』)，便是通过PPO的方式去微调LM，需要的数据便是三元组「query, response, reward」，具体如下图所示

1. Rollout：语言模型根据query生成response
2. 评估：怎么评估模型针对特定query生成response的质量呢，我们可以使用a function、model、human feedback或它们的某种组合进行评估，然后为每个query/response对产生一个标量值，说白了 就是奖励模型有了，那就直接打分
3. 优化：在优化步骤中，「query/response pairs」用于计算序列中标记的对数概率，且比较下面这两个模型输出之间的 KL 散度用作额外的奖励信号
     经过训练的模型(即上图中的Active model)
     基线模型(即上图中的Reference model)，通常是PPO微调之前的模型(比如这里的GPT2，或者instructGPT里的SFT)
   最终，使得Active model生成的响应不会偏离基线模型Reference model太远

PPO算法是一种具体的Actor-Critic算法实现，比如在对话机器人中，输入的prompt是state，输出的response是action，想要得到的策略就是怎么从prompt生成action能够得到最大的reward，也就是拟合人类的偏好。具体实现时，可以按如下两大步骤实现

1. 首先定义4个模型：Actor(action_logits)、SFT(sft_logits)、Critic(value)、RM「r(x, y)」，和kl_div、reward、优势函数adv
   从prompt库中采样出来的prompt在经过SFT(微调过GPT3/GPT3.5的模型称之为SFT)做generate得到一个response，这个『prompt + response』定义为sequence(这个采样的过程是批量采样进行generate，得到一个sequence buffer)，然后这个sequence buffer的内容做batched之后输入给4个模型做inference 这4个模型分别为Actor、SFT、Critic、RM，其中：
   Actor和SFT都是175B的模型，且Actor参数由SFT初始化(SFT是baseline)，Actor输出action_logits，SFT输出sft_logits
   sft_logits和action_logits做kl_div，为了约束actor模型的更新step不要偏离原始模型SFT太远
   
   Critic和RM是6B的模型，Critic参数由RM初始化
   Critic输出标量value，RM输出标量r(x, y)，由r(x, y)和kl_div计算得到reward，reward和value计算得到adv
2. 其次，通过pg_loss和value_loss优化迭代
   Actor的流程是取出sequence，然后inference生成新的logits，再和sequence对应的之前的logits计算ratio，和adv计算出pg_loss，也就是actor的loss，然后反向传播，优化器迭代
   Critic的流程是取出sequence，然后inference得到新的value，和old_value做clip_value，再和reward计算value loss，然后反向传播，优化器迭代

以下是计算策略损失和价值损失的关键代码(来自trl/ppo_trainer.py at main · lvwerra/trl · GitHub的第971-1032行)，且为方便大家阅读时一目了然，我特意给每一行的代码都加上了注释

上面代码中 有两点值得解释下

1. 计算回报时，为何是优势值加上对应的value值
   别忘了，根据本博客中另一篇文章《RL极简入门》可知
   优势函数A(s,a)定义为Q(s,a) - V(s)，其中Q(s,a)是动作价值函数，表示在状态s采取动作a所能获得的预期回报
   而V(s)则是状态价值函数，表示在状态s下依据当前策略所能获得的预期回报

   其训练过程类似 ChatGPT，而通过本博客内的《ChatGPT技术原理解析》3.1节，可知训练三个模型(SFT、RM、RL/PPO)得先准备三套数据集

   2.1.1 actor_training_data，即用于微调GPT3所用的数据

   actor_training_data，即用于微调GPT3所用的数据，比如
   [
     {
         "user_input": "here the input of the user",
         "completion": "here the model completion"
     }
   ]
   actor_training_data如何而来呢，有4项途径

   1. 使用 100% 合成数据，可以通过运行以下命令综合生成数据集：
      python artifacts/generate_actor_dataset.py，注：此命令需要订阅OpenAI，生成完整数据集的davinci-003成本约为 200 美元(当然 也有免费的途径)
   2. 使用具有辅助交互的开源数据集之一，目前支持：

      Anthropic HH RLHF：这个数据集由结构化的 {question/answer pairs} 组成，包括机器人选择和拒绝的答案；
      Stanford Human Preferences Dataset (SHP)：这个数据集是从选定的“提问”subreddits 中挑选出来的，并且包括基于最受支持的回答的范围广泛的 {question/answer pairs} 的问题
      可以运行以下命令下载数据集：

      其中：
      <dataset_name>对于 StanfordNLP/SHP 数据集，可以是“SHP”或“ARLHF”，对于 Anthropic/hh-rlhf 数据集，可以分别是“SHP”或“ARLHF”；
      <path_to_folder_for_download>是要创建数据集的文件夹路径；
      <N>是组成 reward_dataset.json 的样本数

   3. 使用 100% 个性化数据集
      用户提供自己的个性化完整数据集，数据集必须是具有以下格式的 JSON 文件：
      [
          {
              "user_input": "here the input of the user",
              "completion": "here the model completion"
          }
      ]
      其中列表包含多个dictionaries，每个dictionary 对应一个数据样本，建议使用超过 1000 个数据样本来进行对actor的训练

   4. 创建完整的数据集，增加一些自定义数据样本，数据集可以从用户提供的一些提示+响应示例中综合生成（少数 => 10）

   2.1.2 reward_training_data，用于训练一个奖励模型的数据

   reward_training_data，用于训练一个奖励模型的数据，包含三部分的数据： 
   i) prompts,
   ii) completion
   iii) score of the completion assigned accordingly to the user feedback (the Human Feedback in RLHF，即对各个回答的评分score)
   
   示例如下
   [{
       "user_input": "...",
       "completion": "...",
       "score": 1
   },
       ...
   ]
   
   同样的，奖励数据怎么来呢？有以下三种方式

   1. be synthetically scored using a LLM as Human Feedback
      LLM 模型用于为每个entry计算分数
      为此，LLM 需要一个提示模板，其中包含评估生成的文本的所有说明(比如奖励规则，什么情况下该奖 什么情况下不奖都得十分明确)。为此，您应该将key reward添加到文件中templates.json，比如：

      {

          "reward": "Here is the template for the reward model. The rules are: 1.Rule 1 2. Rule 2"
      }
      如果未提供模板，则使用默认模板artifacts/generate_rewards.py，注：所有模板都必须保存在一个名为 .json 的 JSON 文件中templates.json

      获得unlabelled dataset后，您可以通过运行以下命令生成分数：

      其中，<dataset_path>要评分的reward dataset的路径；
      <model_to_use>用于奖励的模型，默认建议使用text-davinci-003
      <temperature>用于对模型进行评分的temperature，temperature =0.1；
      <max_tokens>
      <reward_template>，这是包含用于生成奖励的模板的文件的路径，如果未提供路径，将使用默认模版
      
      这里值得注意的是，与instructGPT中的「人类通过对模型的输出进行排序，然后利用这些排序数据去训练一个RM」不同，ChatLLaMA直接训练一个RM对模型的输出进行打分 比如0-5分，且与人类的打分做MSE损失(减少RM打分与人类打分之间的差距)

      最后，可能你会问，从哪里看出来的用的MSE损失，答案是从另外一个文件里看出来的(具体是chatllama/rlhf/reward.py 文件的第282行) 

   2. 用户提供他们个性化的完整数据集(至少需要 100 个数据样本)，但数据集必须是以下格式的 JSON 文件，取名为：reward_training_data.json

   3. 用户提供的少量示例和使用 LLM 综合扩展的数据集(通过self-instruct的方式提示LLM产生更多所需要的指令数据)

   2.1.3 rlhf_training_data，通过RL方法不断优化迭代最优策略的数据

   It can be provided in 2 different ways:

   • Few examples provided by the user and dataset synthetically expanded using LLM(依然可以

     继续通过self-instruct的方式提示LLM产生更多所需要的指令数据)
     需要将key rlhf添加到templates.json文件中，其中包含有关要执行的任务的信息以及 LLM 生成所需的额外上下文，这是模板的示例(所有模板必须保存在一个名为templates.json)：

     {
       "rlhf": "Here is the template for the generating RLHF prompts. The task we want to perform is ..."
     }

   • The user provides the full dataset with possible interactions with the model

     数据集需要包含超过 1000 个提示示例(文件命名为rlhf_training_data.json)：

     [
         {
             "user_input": "here the example of user input"
         }
     ]

   2.2.1 RewardTrainer 类的 train 方法训练一个奖励函数

   chatllama/rlhf/reward.py中

   首先定义了一个名为 Reward Model 的类，作为奖励模型或批评者模型(Critic Model)。Reward Model 是一个基于语言模型的模型，附加了一个头部head，用于预测给定的 token 序列的奖励(一个标量值)，最后将CriticModel类设置为RewardModel类，以保持命名一致性

   之后，定义类：RewardDatase用于训练奖励模型的数据集
   RewardDataset 类是一个继承自 Dataset 的自定义数据集类，它的作用是从给定的 JSON 文件中读取数据，并将数据整理成适当的格式。JSON 文件应包含以下格式的数据：

   其中 user_input 是用户的初始输入，completion 是模型生成的补全，而 score 是用户或LLM给予补全的分数

   再定义一个RewardTrainer 类用于训练奖励模型，它初始化奖励模型、优化器、损失函数(具体如上文所说，或如282行所述的MSE损失函数)、数据集和数据加载器等。此外，它还支持使用 DeepSpeed 或 Accelerate（两种高性能深度学习训练框架）进行训练
   
   RewardTrainer 类的主要方法有：
   train：训练奖励模型。它执行训练循环，包括前向传播、计算损失、反向传播和优化器更新。在每个周期结束时，它还可以对模型进行验证（如果提供了验证数据集的话）

   • 首先是初始化

   • save_checkpoint：保存模型的检查点。在训练过程中，可以定期将模型的当前状态（包括模型参数、优化器状态和训练统计信息）保存为检查点，以便在需要时恢复训练
   • load_checkpoint：从检查点恢复模型。如果在训练过程中找到检查点文件，则该方法将从检查点恢复模型状态，并返回从何处恢复训练的周期和步骤
   • 接下来，是具体的训练过程

   总之，在 RewardTrainer 类的 train 方法中
   首先会尝试从检查点恢复模型（如果有的话）；
   然后，它会遍历数据加载器中的所有输入，对每个输入执行前向传播、计算损失、反向传播和优化器更新；在每个周期结束时，如果提供了验证数据集，还会对模型进行验证；
   最后，在训练完成后，将保存模型

   2.2.2 通过chatllama/rlhf/actor.py再训练一个actor

   此外，项目通过chatllama/rlhf/actor.py再训练一个actor，比如通过train方法训练一个基于transformer的模型，它包括了数据处理、模型训练、验证和模型保存等操作

   1. 定义方法，它没有返回值。
   2. 打印训练开始信息。
   3. 获取配置参数，包括批量大小、训练轮数、设备和检查点步数。
   4. 计算迭代次数。
   5. 加载模型检查点并获取开始的轮数和步数。
   6. 如果从头开始训练，清空训练统计。
   7. 初始化检查点计数器。
   8. 定义训练循环，其中包括：
      • 设置模型为训练模式。
      • 遍历训练数据加载器。
      • 如果恢复训练，跳过已经完成的步数。
      • 对输入文本进行标记化处理。
      • 将输入文本分割成tokens和mask。
      • 添加结束符（EOS）。
      • 将输入文本分割成输入和输出。
      • 将输入文本移至设备。
      • 执行前向传播。
      • 计算损失。
      • 执行反向传播和优化。
      • 打印训练进度。
      • 定期保存检查点和训练统计。
   9. 进行验证（如果启用验证的话）：
   10. 设置模型为评估模式。
   11. 使用禁用梯度计算。
   12. 遍历验证数据加载器。
   13. 对输入文本进行标记化处理。
   14. 将输入文本分割成验证输入和输出。
   15. 执行前向传播。
   16. 计算损失。
   17. 更新验证损失统计。
   18. 打印验证进度。
   19. 在恢复训练后，将重置为0。
   20. 训练完成后，保存模型。
   21. 打印训练结束信息

   2.2.3 通过PPO算法优化强化学习任务中的策略(actor)和价值(critic)网络

   有了奖励函数和actor，便可以通过PPO算法优化强化学习任务中的策略(actor)和价值(critic)网络，具体如下图，设置内外两个循环

   • 外层循环迭代训练轮次(epochs)
   • 内层循环遍历数据加载器(dataloader)中的批次(batches)，在每次迭代中，它会处理一批数据，包括状态、动作、价值等，这些数据用于训练智能体-评论家模型

   在内层循环中依次做如下处理(以下代码来源于：chatllama/chatllama/rlhf/trainer.py )：

   首先是导入必须的库和模块，当然，主要是

   1. 导入所需的库和模块。
   2. 函数：用于在两个不同的分词器之间转换给定的tokens。
   3. 函数：检查两个配置是否属于相同的模型家族。
   4. 类：包含了actor和critic模型，并用于在训练actor过程中为给定的序列生成动作和值。它包括以下方法：
      • ：初始化actor和critic模型

      • ：加载模型，但未实现。
      • ：将模型保存到路径。
      • ：基于给定的整个序列，使用actor的forward方法获取序列中每个token的logits，并使用critic的forward方法获取每个生成步骤的值。

   这个代码主要用于强化学习训练自然语言生成模型。类是其中的核心部分，它包含了actor和critic模型。这两个模型在训练过程中相互协作，用于生成动作和值。

   其次，主要是关于一个用于生成动作、动作逻辑、价值和序列的生成函数，以及用于存储训练数据和生成训练示例的类

   1. 首先定义了一个名为generate的函数，它使用了@torch.no_grad()和@beartype修饰器。这个函数接收四个参数，分别是states_actor、states_mask_actor和states_critic，并返回一个元组
      这个函数的主要目的是从输入状态生成动作、动作逻辑、价值和序列。它首先从actor模型生成动作序列，然后创建一个用于actor序列的mask。接下来，它检查是否需要为critic使用不同的编码。如果需要，它将使用change_tokenization函数来更改序列的编码。接着，它生成动作逻辑和价值。如果处于调试模式，将打印一些调试信息。
   2. 接下来，代码定义了一个名为Memory的namedtuple，用于存储每个经验的数据。Memory包含了11个字段，如states_actor、actions、values等。
   3. 然后定义了一个名为ExperienceDataset的类，它继承自torch.utils.data.Dataset。这个类用于训练actor-critic模型。它接收一个memories参数和一个device参数。memories参数是一个包含Memory实例的双端队列。device参数表示要在哪个设备上进行计算。这个类实现了__len__和__getitem__方法，使其可以像普通的PyTorch数据集一样使用。
   4. 最后，定义了一个名为ExamplesSampler的类，用于从JSON文件中读取示例并在需要时抽样。这个类接收一个表示文件路径的path参数。在初始化时，它从文件中读取数据，并将其存储在self.data中。它还实现了一个名为sample的方法，用于从数据中抽取指定数量的示例。

   再之后，定义了一个名为 的类，用于使用强化学习训练一个Actor-Critic模型。该类具有多个属性和方法，用于训练过程中的各种操作。

   • 在 方法中，初始化了训练器的各个组件，包括Actor-Critic模型、actor和critic优化器、reward模型、用于存储训练统计数据和对话记录的类、以及示例采样器
   • 方法保存了当前状态的Actor-Critic模型的检查点，包括当前的训练轮数、actor和critic模型的状态字典，以及它们各自的优化器的状态字典。
   • 方法加载了Actor-Critic模型的检查点，包括训练轮数、actor和critic模型的状态字典，以及它们各自的优化器的状态字典。如果没有找到检查点，则返回轮数0。如果actor和critic的检查点存在差异，则从两者中最小的轮数开始训练。

   再之后，调用 方法更新actor和critic模型，并保存训练统计数据和对话记录

   1. 使用智能体-评论家模型计算新的动作概率和价值

   2. 计算动作的对数概率、熵和KL散度损失

   3. 计算重要性权重比率(ratios)，即新旧策略的概率比

   4. 计算PPO损失，包括优势函数的计算和PPO-clip算法的应用
      首先我们回顾下强化学习极简入门一文里对『近端策略优化裁剪PPO-clip』的阐述 简单实现的话，即是
      更具体的实现，则可以如下所示

   5. 计算策略损失和总损失
      可能有读者看到这里 看迷糊了，即咋出来两个损失函数了，看起来是一个策略损失，一个KL散度损失，与我们在本博客里的另一篇文章《》中「3.1.3 InstructGPT训练阶段3：如何通过PPO算法进一步优化模型的策略」探讨的结果咋不太一样呢 ？ 不急，我们先来分析下这两个损失函数
        一个 policy loss，本质是一个目标函数(具体用的近端策略优化裁剪PPO-clip与熵的差值) 其中，表示在当前策略下采样得到的经验的无偏估计， 是策略比率，是优势函数， 是超参数，用于控制策略更新的幅度，是熵的系数
      
      当然 在instructGPT的原理中并没有这个熵
        另一个 KL散度损失(kl_div_loss)，还是用于限制新策略与旧策略之间的差异，以免更新太快，导致学习不稳定

      对应的公式为

      值得一提的是，这里确实容易引发疑惑，毕竟上面的policy loss已经对新旧策略的比值 ratios = (actions_log_prob - old_actions_log_probs).exp() 做了截断处理，而这里又加一个对新旧策略差值的KL散度约束，未免有多此一举之嫌，比如在instructGPT的原理中便只有两者其一：关于策略梯度的损失就一个policy loss

      最终，总的损失函数为：

      其中， 是超参数，用于控制 KL 散度损失的权重

   6. 如果损失值为NaN，抛出异常

   7. 更新策略，包括使用DeepSpeed或Accelerate库进行优化

   8. 计算价值损失
      本文发布后，有读者留言对这块表达疑惑，即怎么是先计算裁剪的损失，然后对比未裁剪的损失，然后两种损失中取更大呢，原因和上文第六部分最后解释的一样，便不再重复了
   9. 如果价值损失为NaN，抛出异常

   10. 更新评论家，包括使用DeepSpeed或Accelerate库进行优化

   11. 将损失值追加到训练统计信息中

   12. 输出迭代信息

   13. 训练循环结束后，将智能体-评论家模型设为评估模式并输出训练结束信息

   最后的最后，定义了一个 train() 方法，使用 actor-critic 算法训练强化学习模型。方法首先初始化各种设置，如训练的总 episode 数量、每个 episode 的最大步数、批次大小和训练设备等。然后检查要用于学习的记忆数量是否是批次大小的倍数，以及总步数是否是更新步数的倍数。

   该方法初始化记忆，加载检查点（如果有的话），如果是从头开始的新训练，则清除会话记录。然后循环遍历 episode 和 timestep，从示例数据集中抽取样本，为 actor 和 critic 进行分词，生成动作和值的序列，计算动作日志概率，计算奖励。存储每个 episode/timestep 的记忆，并将完成（解码后的动作）记录在会话日志中。

   在一定数量的 timestep 后，使用记忆进行学习，并计算平均奖励。该过程重复进行，直到训练完成。该方法在训练结束时保存模型和会话日志。

   ---

   据介绍(介绍页面，该页面的翻译之一，代码地址)，Colossal-AI 开源了基于 LLaMA-7B 模型的包含完整 RLHF 流程的类 Chat 模型复现方案 ColossalChat

   3.1.1 针对社交平台的种子数据且利用self-instruct 技术生成中英双语数据集

   ColossalChat 收集并清洗了社交平台上人们的真实提问场景作为种子数据集，然后利用 self-instruct 技术扩充数据(通过prompt OpenAI API，花费约 900 美元进行标注)，最终生成了10.4万条问答的中、英双语数据集(这是数据的开源地址)
   他们的说法是，对比其他 self-instruct 方法生成的数据集，该数据集的种子数据更加真实、丰富，生成的数据集涵盖的话题更多，该数据可以同时用于微调和 RLHF 训练，通过高质量的数据，ColossalChat 能进行更好地对话交互，同时支持中文

   3.1.2​ ColossalChat训练方式：类似instructGPT/ChatGPT的训练三步骤

   关于训练方式：类似instructGPT/ChatGPT的训练三步骤(如果忘了，务必复习下此文的3.1节)

   • Stage1 是supervised-fintuning，即使用上文提到的数据集进行监督微调
   • Stage2 训练一个奖励模型(初始化为阶段1的SFT模型)，它通过模型对于同一个 prompt 的不同输出进行人工排序，根据排序结果监督训练出一个奖励模型
   • Stage3 是通过阶段2训练出来的奖励函数微调出一个RL模型，微调过程中通过PPO算法限制RL模型的参数更新范围(以阶段1的SFT模型的策略为参考基准，PPO算法避免与基线模型SFT的策略偏离过远)
   • 如上图底部，首先是 Make Experience 部分，利用 SFT 、Actor、RM、Critic模型计算生成 Experience 存入 buffer 中； 之后是参数更新部分，利用 Experience 计算价值损失(value loss)​和策略损失(policy loss)，具体说明在有介绍
   • 如上图顶部即是PTX 部分(上面的目标函数​中加在最后的偏置项)
     ColossalChat 计算 Actor 的现有输出response 和预训练语料的回答部分的交叉熵损失函数(calculates the cross-entropy loss between the Actor’s output response and the response part of the input corpus)
     用来在 PPO 梯度中加入预训练梯度(add pre-training gradients to the PPO gradient)
     以保持语言模型比如GPT2原有的核心性能(maintain the language model’s original performance and prevent forgetting)，防止忘了最早从哪里出发的(GPT2 ​ SFT ​ RM ​ RLHF)
   • 最后将策略损失、价值损失和 PTX 损失加和(the policy loss, value loss, and PTX loss are summed up)，进行反向传播和参数更新 

   先看下整体的代码架构图

   3.2.1 首先，训练一个SFT模型

   首先通过ColossalAI/applications/Chat/coati/trainer/sft.py，训练一个SFT模型

   
   

   3.2.2 训练一个奖励模型

   其次，通过ColossalAI/applications/Chat/coati/trainer/rm.py 训练一个奖励模型

   3.2.3 通过trainer/ppo.py to start PPO training

   最后，通过ColossalAI/applications/Chat/coati/trainer/ppo.py to start PPO training

   在获得最终模型权重后，还可通过量化降低推理硬件成本，并启动在线推理服务，仅需单张约 4GB 显存的 GPU 即可完成 70 亿参数模型推理服务部署

   ---