从零实现带RLHF的类ChatGPT：逐行解析微软DeepSpeed Chat

写在最前面

本文最早写于2023年4月的这篇文章中《从零实现带RLHF的类ChatGPT：从TRL/ChatLLaMA/ColossalChat到DeepSpeed Chat》，后因要在「大模型项目开发线下营」上讲DSC的实现而不断扩写其中的DSC，为避免原文过长，故把该文最后的DSC部分抽取出来成本文

前言

如此文所述，微软开源的DeepSpeed Chat(简称DSC)实现的不错，其类似ChatGPT的三阶段训练方式，会给你一个完整而通透的“PPO算法/RLHF”的代码实现全流程(好的资料可以让你事半功)

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

• 其论文地址：DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales
  其官方的GitHub介绍页面：中文版、英文版
• 其GitHub代码地址：https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat

DSC充分利用DeepSpeed项目的优化方案，考虑actor在RLHF阶段要在train (参数更新)和eval(经验采集)模式中反复切换，不做优化的话整体速度很慢，但是原本DeepSpeed的train加速和eval加速属于是解离的两种方案，DSC就设计了一种叫做DeepSpeedHybridEngine的引擎，使得actor在RLHF阶段能同时享有train和eval加速优化，整体提高RLHF速度

 一句话总结就是：DeepSpeed来给RLHF提速，遂成deepspeed chat

注：七月在线ChatGPT课的一学员“吹牛班的春天”把这个模型写的很细致了(年初至今的5个多月下来，除了本博客内的ChatGPT系列，春天这个deepspeed chat解析是我个人看到的唯一足够深入、细致的，主要真正写的深入、细致的文章实在是太少了，一方面 技术太新，二方面 涉及的细节太多)，故本文大部分的分析基于他的博客修改得到

总的来说，DeepSpeed Chat和instructGPT的三阶段训练方式差不多，该三阶段分别用phase1、phase2、phase3表示

下面简述这训练的三大阶段

第一部分 DSC之phase-1: Supervised Finetuning

1.1 SFT的训练流程

phase1的核心代码见：applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py，至于其训练过程如下图所示( 鼠标右键点击图片：在新标签页中打开图片，可以查看高清大图 )

1. 载入tokenizer(1-2)
2. 载入基座模型（目前仅支持部分CausalLM模型）(3-4)
3. 根据是否设置lora_dim（LoRA的低秩维度）判断是否启用LoRA技术， 如果启用，则将基座模型结构进行LoRA改造（具体可见后续详述），并返回改造后的模型(5-6)
4. 判断是否启用“仅更新LoRA参数”，如果启用，则对其余结构参数进行冻结处理，并返回冻结处理后的模型(7-8)
5. 获取Dataset (9-10)
6. 实例化DataLoader(11)
7. 使用DeepSpeed的优化技术DeepSpeedEngine包裹模型等对象(12)
8. 开始正式训练前首先进行指标评估，选用的指标为困惑度perplexity(13-14)
9. 开始训练，epoch循环：
   

1.2 关于LoRA与困惑度的说明

上述过程有2个细节，值得一提

1. 关于LoRA的详解，可看此文《Alpaca-LoRA：通过PEFT库在消费级GPU上微调「基于LLaMA的Alpaca」》的2.2.3节
2. DeepSpeed-Chat选择了困惑度perplexity作为phase1训练期间的评估指标
   困惑度perplexity是一种度量语言模型性能的指标，它衡量了训练好的模型对测试数据的拟合程度，对于输出句子的每个token，都可以得到其输出的置信概率值，将这些值相乘并取其几何平均数的倒数即可计算得到困惑度perplexity，使用公式表达更为简洁：
   
   其中，输出的句子共有个token，第 个token的置信概率值为
   
   而CausalLM模型的训练过程通常采用对数似然损失来进行优化，其输出的损失公式如下：
   
   其中，输出的句子共有个token，第t tt个token的置信概率值为
   
   因此perplexity与CausalLM的loss之间实际存在如下关系：
   
   相关源码的perplexity计算也是基于上述公式得到的：先是将验证数据输入至模型，得到模型loss输出，然后通过perplexity与loss之间的指数关系计算得到perplexity

       def evaluation(model, eval_dataloader):
           """
           以困惑度perplexity为评估指标进行验证
           """
           model.eval()
           losses = 0
           for step, batch in enumerate(eval_dataloader):
               """
               batch: 由input_ids、attention_mask、labels共3个部分组成的dict。
               其中每个部分的shape均为(bs, max_seq_len)
               """
               batch = to_device(batch, device)
               with torch.no_grad():
                   outputs = model(**batch)
   
               """Causal LM 的损失函数为交叉熵损失"""
               loss = outputs.loss
               losses += loss.float()
           losses = losses / (step + 1)
   
           try:
               """困惑度perplexity通常可以通过exp(CELoss)计算得到"""
               perplexity = torch.exp(losses)
           except OverflowError:
               perplexity = float("inf")
   
           try:
           	"""
           	- get_all_reduce_mean中调用了torch.distributed.all_reduce(perplexity, op=torch.distributed.ReduceOp.SUM)
           	- 对所有进程、或者说GPU（因为通常情况下就是单个进程控制单个GPU）中的perplexity进行求和
           	- 然后再除以全局进程数torch.distributed.get_world_size()得到平均的perplexity结果
           	"""
               perplexity = get_all_reduce_mean(perplexity).item()
           except:
               pass
           return perplexity

第二部分 DSC之phase-2: Reward Model Finetuning

2.1 训练数据

模型将基于排序损失对形如下述样例的数据对进行训练，最后将得到具备类人评分能力的RM(Reward Model)

数据格式名称	说明	样例
chosen_sentence	人类偏好的完整对话，由prompt衔接偏好应答chosen得到，适用于phase1和phase2	“Human: Please tell me about Microsoft in a few sentence? Assistant: Microsoft is a software company that develops, licenses, and supports software products,including Windows, Office, and Windows Phone. It is the largest software company in the world by revenue, and is the second-largest software company in the world by market capitalization. Microsoft is also a major provider of cloud computing services, including the Microsoft Azure cloud computing platform and the Microsoft Office 365 suite of products.”
reject_sentence	人类排斥的完整对话，由prompt衔接排斥应答rejected得到，适用于phase2	“Human: Please tell me about Microsoft in a few sentence? Assistant: I’m not sure what you mean.”

2.2 训练流程

phase2的大致训练过程如UML时序图所示(鼠标右键点击图片：在新标签页中打开图片，可以查看高清大图)：

1. 载入tokenizer(1-2)
2. 载入模型（rm_model），其中涉及一定的结构更改(3-8)
3. 根据是否设置lora_dim（LoRA的低秩维度）判断是否启用LoRA技术， 如果启用，则将基座模型结构进行LoRA改造（具体可见后续详述），并返回改造后的模型(9-10)
4. 判断是否启用“仅更新LoRA参数”，如果启用，则对其余结构参数进行冻结处理，并返回冻结处理后的模型(11-12)
5. 获取Dataset(13-14)
6. 实例化DataCollator，用于进一步对加载的数据进行整理(15-16)
7. 实例化DataLoader(17)
8. 使用DeepSpeed的优化技术DeepSpeedEngine包裹rm_model等对象(18)
9. 开始正式训练前首先进行指标评估，选用的指标为排序结果的准确率accuracy(19-20)
10. 开始训练，epoch循环：
    

2.3 关键代码详解

2.3.1 RM的具体结构

首先使用transformers的AutoModel类来读取指定模型的主干网络(不直接定义有输出头的网络结构)，然后引入一个可实现从hidden_size降维至1的线性层，该线性层将作为主干网络的输出头，为输入序列的每个位置输出1个评分

# applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
"""
rm_model调用了create_critic_model进行载入
默认情况下rm_model是不启用dropout的
"""
rm_model = create_critic_model(···)

# applications/DeepSpeed-Chat/training/utils/model/model_utils.py
def create_critic_model(···):
    """此处的模型读取方法用的是“AutoModel”，因此此处critic_model只有主干部分"""
    critic_model = create_hf_model(AutoModel, ···)

    """
    critic_model传入RewardModel，将额外得到线性层输出头，
    因此此处的critic_model结构为“v_head + 主干部分”
	"""
    critic_model = RewardModel(critic_model, ···)
    ...
    return critic_model

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):
    """
    将读取得到的model的结构修改为适用于RewardModel的形式，
    总的来说即是使用载入的主干网络进行特征提取，
    其所提取的特征（最后层的各位置输出特征hidden_states）将被传入线性层，输出得到1个数值，
    该数值即为分值，因此max_seq_len维度的每个位置均会得到1个分值
    """
    def __init__(self, base_model, ...):
        super().__init__()
		···
        if hasattr(self.config, "word_embed_proj_dim"):
        	"""
			OPT系列模型的word_embed_proj_dim为embedding层的输出维度，
			通常在transformer模型中也就等于 hidden_size，
			v_head将基于主干网络的输出特征 hidden_state 进行分值预测，共输出max_seq_len个分值
			"""
            self.v_head = nn.Linear(self.config.word_embed_proj_dim,
                                    1,
                                    bias=False)
        ···
        """base_model即为主干网络，因此RM最终由1个主干网络和1个线性层构成"""
        self.rwtranrsformer = base_model

RM的模型结构基本如下所示(此处的基座模型为“facebook/opt-125m”)，由主干网络rwtransformer及输出头v_head组成：

RewardModel(
  (v_head): Linear(in_features=768, out_features=1, bias=False)
  (rwtranrsformer): OPTModel(
    (decoder): OPTDecoder(
      (embed_tokens): Embedding(50272, 768, padding_idx=1)
      (embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
      (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (layers): ModuleList(
        (0-11): 12 x OPTDecoderLayer(
          (self_attn): OPTAttention(
            (k_proj): Linear(in_features=768, out_features=768, bias=True)
            (v_proj): Linear(in_features=768, out_features=768, bias=True)
            (q_proj): Linear(in_features=768, out_features=768, bias=True)
            (out_proj): Linear(in_features=768, out_features=768, bias=True)
          )
          (activation_fn): ReLU()
          (self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (fc1): Linear(in_features=768, out_features=3072, bias=True)
          (fc2): Linear(in_features=3072, out_features=768, bias=True)
          (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        )
      )
    )
  )
)

2.3.2 DataCollator及RM所需输入形式

phase2使用的数据整理器data_collator为DataCollatorReward()，本阶段取出的单个样本example实际上是一个chosen-rejected数据对（见下方代码块）

即1个大小为batch_size的batch取出了batch_size个数据对，data_collator将把数据对拆成chosen_sentence和reject_sentence（example一分为二），因此实际上1个batch真正输入模型的数据量大小应当为“batch_size * 2”

# applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
"""phase2使用的data_collator为DataCollatorReward()"""
data_collator = DataCollatorReward()

# applications/DeepSpeed-Chat/training/utils/data/data_utils.py
class DataCollatorReward:
    def __call__(self, data):
        """
        对dataloader取到的数据 data 进一步整理，将数据整理成batch输入形式
        入参 data 的具体样式可见下个代码块
        """
        batch = {}

        """f为data中的1个tuple，tuple的第0个元素和第2个元素
        分别为chosen_sentence和reject_sentence的input_ids"""
        batch["input_ids"] = torch.cat([f[0] for f in data] + 
        							   [f[2] for f in data],
                                       dim=0)

        """f为data中的1个tuple，tuple的第1个元素和第3个元素
        分别为chosen_sentence和reject_sentence的attention_mask"""
        batch["attention_mask"] = torch.cat([f[1] for f in data] +
                                            [f[3] for f in data],
                                            dim=0)

        """batch的具体样式可见下个代码块"""
        return batch

且输入的data为一个batch的数据列表，其中的 每个元素 为一对chosen-rejected数据：

    (
	 chosen_sentence_input_ids, 
	 chosen_sentence_attention_mask,
	 reject_sentence_input_ids,
	 reject_sentence_attention_mask
	)

每组数据的第0个元素和第2个元素为input_ids，第1个元素和第3个元素为attention_mask

输出的batch为字典：{“input_ids”: tensor([...]), "attention_mask": tensor([...])}
并且字典值中chosen位于前半部分，rejected位于后半部分：

    {
    "input_ids": [
                  chosen_sentence_1_input_ids,
                  chosen_sentence_2_input_ids,
                  ...,
                  reject_sentence_1_input_ids,
                  reject_sentence_2_input_ids,
                  ...
                 ]
    "attention_mask": [
                       chosen_sentence_1_attention_mask,
                       chosen_sentence_2_attention_mask,
                       ...,
                       reject_sentence_1_attention_mask,
                       reject_sentence_2_attention_mask,
                       ...
                      ]
        
    }

后续输入模型后，直接将数据切分出前半部分和后半部分进行并列，即可获得对应的chosen-rejected数据对

2.3.3 整个对话的reward设计和成对排序损失

RM的正向传播过程不算复杂，总的来说就是：

1. 数据经过主干网络得到shape为(bs*2, max_seq_len, hidden_size)的最后层输出特征hidden_states；
2. 然后将输出特征送入线性层v_head得到shape为(bs*2, max_seq_len)的评分rewards

较为复杂的部分实际上是“成对排序损失的计算”以及“评分聚合设计”

2.3.3.1 成对排序损失（Pairwise Ranking Loss）

其中，为RM，为prompt，为chosen，为rejected，和则分别为chosen_sentence和reject_sentence。
该损失函数的目的在于最大化“chosen/好的/排序靠前的”和“rejected/坏的/排序靠后的”的差值，由此促使学习到相应的排序模式

DeepSpeed-Chat在实现这部分时，和分别选择了chosen_sentence和reject_sentence两者answer的对齐部分，通过文字叙述略显抽象，查看下方的代码块有助于你理解这个概念：

max_seq_len为10，pad_token_id为0，
有同属同个prompt的chosen_sentence和reject_sentence:
prompt: [11, 22, 33]
chosen_sentence: [11, 22, 33, 44, 55, 66, 0, 0, 0, 0]
reject_sentence: [11, 22, 33, 40, 50, 0, 0, 0, 0, 0]

“两者answer的对齐部分”即为“非prompt部分也非padding部分、但长度要对齐”：
chosen_truncated: [44, 55, 66]
reject_truncated: [40, 50, 0]

所以当上面的chosen_sentence的answer比较长时，reject_sentence在取相应部分时要取至与chosen部分等长为止；
类似的，如果reject_sentence的answer较长时，同理

为了取到上述提及的“对齐部分”，代码进行了较为晦涩抽象的取index操作，但只要理解其最终目的是为了取到chosen_sentence和reject_sentence对齐部分的reward，来进行损失计算即可

2.3.3.2 对话奖励设计

        尽管使用的是“对齐部分”的reward来计算成对排序损失，但RM模型对一个对话的预测评分实际上取的是该对话文本最后一个有效token(通常会是“结束标记”)的reward，下方代码块提供了一个简单例子说明了这个情况

pad_token_id = 0
conversation = [11, 22, 33, 44, 55, 66, 0, 0, 0, 0]
conversation_rewards = [2.01, 0.23, 2.89, 0.66, 0.33, 2.25, 0.36, 0.99, 1.32, 1.62]
token_id为66的token作为该对话的最后1个有效token，
其对应的reward“2.25”将被用于表示整个对话的reward

整体代码如下所示

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):
	def __init__(self, ···):
		···
	···
	def forward(self, input_ids=None, ···):
		"""获得主干网络的输出的特征"""
		transformer_outputs = self.rwtranrsformer(···)

		"""
		取最后一层的输出特征
		hidden_states.shape: (bs*2, max_seq_len, hidden_size)
		"""
		hidden_states = transformer_outputs[0]

		"""
		将特征送入全连接层得到分数回归值
		rewards.shape: (bs*2, max_seq_len)
		"""
		rewards = self.v_head(hidden_states).squeeze(-1)

		"""先前提及过，实际的bs应该是输入bs的一半"""
		bs = input_ids.shape[0] // 2

		"""区分出chosen和reject"""
		chosen_ids = input_ids[:bs]
		rejected_ids = input_ids[bs:]
		chosen_rewards = rewards[:bs]
		rejected_rewards = rewards[bs:]
		
		loss = 0
        for i in range(bs):
            """
            取出同组chosen和rejected的token_id和分值reward
            chosen_id.shape: (max_seq_len, )
            """
            chosen_id = chosen_ids[i]
            rejected_id = rejected_ids[i]
            chosen_reward = chosen_rewards[i]
            rejected_reward = rejected_rewards[i]

			"""
			下方本应有各种取index相关的操作，
			基于源码解读的可读性考量，且这些部分只是逻辑形式上的弯弯绕绕，与相关原理并不存在直接关系，所以选择暂且将它们忽略
			"""
			
			"""
			c_ind为chosen_sentence的answer后的第一个pad_token的index
			例如pad_token_id=0，sentence[11,22,33,44,55,66,0,0,0,0]，c_ind即为第一个pad_token的index=6 """
            c_ind = ···

            """
            r_ind同理，为reject_sentence的answer后的第一个pad_token的index"""
            r_ind = ···

            """end_ind则为两者的较大者"""
            end_ind = max(c_ind, r_ind)

            # 取chosen和rejected第一个不同的地方的index，可以理解为“response中两个回答自由发挥的第1个token的index”
            """divergence_ind为chosen_sentence和reject_sentence两者answer的第1个token的index"""
            divergence_ind = ···

            """
            以chosen_sentence和reject_sentence最先不同的地方为起始、生成结束的地方为终止，取两者在这个片段的对应分值
            这部分其实就是上个代码块提及的“对齐部分”
            """
            c_truncated_reward = chosen_reward[divergence_ind:end_ind]
            r_truncated_reward = rejected_reward[divergence_ind:end_ind]

            """
            (c_truncated_reward - r_truncated_reward).shape: (truncated_seq_len,)
            计算损失时使用了rank loss的形式，并且是对chosen和rejected“对齐片段”进行计算的
            """
            loss += -torch.log(
                torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
		
        loss = loss / bs
        
        """取代表结束的pad token所在位置的前一个位置（可以理解为的最后一个有效token的位置）的分值作为参考分值"""
            chosen_mean_scores.append(
                chosen_reward[c_ind - 1])  #use the end score for reference
            rejected_mean_scores.append(rejected_reward[r_ind - 1])
        chosen_mean_scores = torch.stack(chosen_mean_scores)
        rejected_mean_scores = torch.stack(rejected_mean_scores)
        
        """返回损失和参考分值"""
        return {
            "loss": loss,
            "chosen_mean_scores": chosen_mean_scores,
            "rejected_mean_scores": rejected_mean_scores,
        }
   ···

2.3.4 phase2的指标评估

DeepSpeed-Chat在phase2中使用的评估指标为排序正确的accuracy，主要过程为：

1. 将数对chosen-rejected数据对（过程中被data_collator拆分为chosen_sentence和reject_sentence）输入RM中进行推理，得到各个sentence的分值；
2. 将同属一个prompt的chosen_sentence得分与reject_sentence得分进行比较，当chosen_sentence得分大于reject_sentence得分时，即为“正确预测”，否则为“错误预测”；
3. 统计正确预测的结果，计算accuracy作为评估指标。
4. 此外评估过程中还将统计平均的chosen_sentence分值“scores”供参考

def evaluation_reward(model, eval_dataloader):
    model.eval()
    """统计预测（赋分）正确的结果
    即 chosen_reward > rejected_reward 的结果数"""
    correct_predictions = 0

    """统计预测总数"""
    total_predictions = 0
    scores = 0
    for step, batch in enumerate(eval_dataloader):
        batch = to_device(batch, device)
        with torch.no_grad():
            """outputs: {'loss':tensor(), 
            			'chosen_mean_scores':tensor(bs,), 
            			'rejected_mean_scores':tensor(bs,)}"""
            outputs = model(**batch)

        """chosen.shape: (bs,)"""
        chosen = outputs["chosen_mean_scores"]

        """rejected.shape: (bs,)"""
        rejected = outputs["rejected_mean_scores"]

        """"赋分正确"即为chosen分值大于rejected分值"""
        correct_predictions += (chosen > rejected).sum()
        total_predictions += chosen.shape[0]

        """累加每个step的平均chosen分值"""
        scores += outputs["chosen_mean_scores"].mean().float()

        if step == 99:  # For faster evaluation and debugging
            break
    """计算acc指标"""
    acc = correct_predictions / total_predictions

    """计算当前step的平均chosen分值"""
    scores = scores / (step + 1)
    try:
        """多进程结果求和求平均"""
        acc = get_all_reduce_mean(acc).item()
        scores = get_all_reduce_mean(scores).item()
    except:
        pass
    return scores, acc

对于RM这块，最后值得一提的是在DeepSpeed-Chat的实现中，RM模型对一个对话的预测评分实际上取的是该对话文本最后一个token的reward，当然此处并不是只能采用这种方式对对话进行评分，这是一个开放性的策略设计，只是DeepSpeed-Chat团队采取了这样的实现，用户当然也可以自己制定评分的处理策略，比如answer部分的平均reward、序列reward再接全连接层得到聚合rewad等等

In our implementation, we use either the end token of the sequence or the first padding token as the aggregated score and compare them. Others may also use the average score for the entire answer as an alternative.

第三部分 DSC之phase-3: RLHF Finetuning

本小节改编自七月在线ChatGPT课学员春天的关于deepspeed chat的解析的第三part

3.1 RLHF的训练数据

数据格式名称	说明	样例
prompt	对当前情境的描述，为模型生成提供指令输入信息，可以理解为通俗含义上的“问句”，适用于phase3	"Human: Please tell me about Microsoft in a few sentence? Assistant: "（举文本例子是为了便于理解，实际上此处为input_ids）
seq	actor基于prompt输入生成的完整对话序列。	"Human: Please tell me about Microsoft in a few sentence? Assistant: Microsoft is a world-renowned company."举文本例子是为了便于理解，实际上此处为input_ids）
logprobs	actor基于seq输出的logits/策略对数	shape: 本应为(seq_bs, max_seq_len, vocab_size)，经过gather处理后仅取实际label token的log_logit值，为(seq_bs, max_seq_len, 1)
ref_logprobs	reference/SFT基于seq输出的logits/策略对数	shape: 本应为(seq_bs, max_seq_len, vocab_size)，经过gather处理后仅取实际label token的log_logit值，为(seq_bs, max_seq_len, 1)
value	critic基于seq输出的对序列每个位置的价值评估	shape: (seq_bs, max_seq_len)
reward	eward/RM基于seq输出的对整个对话的（环境）奖励，实际代码实现时还会再加个β惩罚项	shape: (seq_bs,)
attention_mask	用于滤掉非有效元素	shape: (seq_bs, max_seq_len)

有两点值得重点一提的是

1. 各个框架对于经验数据的定义不完全相同，例如ColossalChat定义的经验数据还比此处多了项“adv”和“reward”（此reward非彼reward，ColossalChat的reward指的是“经过KL散度修正后的KL_Reward”），但本质上都是同理的，只是框定的范围不同，因为adv（优势函数Adventage）和KL_Reward完全可以由已有项logprobs、ref_logprobs、reward、value计算得到
2. 从代码效率的角度来考量，ColossalChat的经验数据定义相对更严谨些，因为优势adv以及KL惩罚奖励完全可以由基本经验数据计算得到，在生成经验的阶段一步到位计算即可
   而DeepSpeed-Chat中将其安排在训练阶段来计算，每次PPO迭代才计算(毕竟优势和KL惩罚奖励是基于基本经验数据计算得到的，而基本经验数据在生成经验阶段已经确定了，所以即使是在不同的PPO迭代中，优势和KL惩罚奖励也是不变的，因此DeepSpeed-Chat对adv以及KL惩罚奖励进行了重复计算，这个环节的计算顺序估计后续相关团队会做出调整)

3.2 RLHF的整个训练过程

整个RLHF的训练过程如下图所示(鼠标右键点击图片：在新标签页中打开图片，可以查看高清大图)

1. 载入tokenizer(1-2)；
2. 获取Dataset并实例化DataCollator(3-9)：获取用于采集经验的prompt的Dataset(4-5)，如果启用了无监督训练，则再获取无监督数据的Dataset(6-7)，并且实例化DataCollator用于进一步对加载的数据进行整理
   
   data_collator实例化自DataCollatorRLHF，该类主要实现了“padding至max_prompt_len (默认为max_seq_len的一半)，然后进行flip”
   为何要专门对prompt token进行flip(翻转)操作？
   原因在于，phase3取用prompt的目的在于，将prompt输入至actor模型中，actor将根据prompt自回归地生成后续内容，以此进行经验采集
   
   以基座为opt-125m的actor模型为例，该模型所能支持的最大序列长度（max_seq_len）为512，而phase3还将预设有最大prompt长度（max_prompt_len），通常为max_seq_len的一半，即256，余下的另一半长度将被用于生成
   那么当输入的prompt不满足最大prompt长度max_prompt_len时，将需要对该prompt进行padding操作（在phase3的data_collator代码中有所体现），而padding操作通常又是直接于序列后侧加入pad token，padding后的输入将变成[prompt, padding]的形式，自回归生成任务将接续pad_token进行生成——这是不合理的
   因此需要先将prompt输入进行flip翻转，翻转后进行padding操作，然后再flip翻转回来，padding后的输入就成为了[padding, prompt]的形式，对于自回归任务来说，接续prompt的内容进行生成才是合理的
   
   通过下述例子应该能更好地理解这个操作的用意

   max_prompt_len = 5
   pad_token_id = 0
   
   prompt_token_ids = [233, 11, 22]
   # padding位于后侧 ×
   prompt_token_ids.padding() = [233, 11, 22, 0, 0]
   
   prompt_token_ids.flip(0) = [22, 11, 233]
   prompt_token_ids.flip(0).padding() = [22, 11, 233, 0, 0]
   # padding位于前侧 √
   prompt_token_ids.flip(0).padding().flip(0) = [0, 0, 233, 11, 22]

3. 实例化DataLoader(10)；
4. 使用DeepSpeedRLHFEngine()载入PPO训练所需的各个模型（actor、ref/SFT、critic、reward/RM），并进行封装得到rlhf_engine(11-12)；
5. 实例化PPO的训练管理trainer(13-14)；
6. 实例化用于PPO训练环节的MiniDataset（有别于上述Dataset，上述Dataset用于获取整个大轮次的数据，MiniDataset进一步管理Dataset提供的数据，用于分配给PPO轮次、即小轮次进行训练）(15-16)；
7. 开始训练，大轮次epoch（prompt_epoch）
   

---

3.3 阶段三的关键代码详解：step3_rlhf_finetuning

3.3.1 阶段三中各个模型的初始化：step3_rlhf_finetuning中的main.py、rlhf_engine.py

关于模型的初始化，源码中使用了DeepSpeedRLHFEngine类进行了actor、ref/SFT、critic、reward/RM、actor_ema等模型的初始化，该类主要实现了：

1. 读取模型，虽然同样支持直接从huggingface hub拉取相应模型，但通常都是从本地路径读取phase1、phase2训练好的模型：
     actor、ref/SFT以及actor_ema (EMA是ExponentialMovingAverage，中文叫做指数移动平均，是一种模型训练技巧，模型在进行第k次更新得到的参数并不是直接使用第k次的新参数，而是由k-1时的历史参数和第k次的新参数加权平均得到，主要是加强训练过程的稳定性 ) 通常都初始化自phase1训练所得的模型；
     critic与reward通常都初始化自phase2训练所得的模型
2. 为各个相关模型设定不同的DeepSpeed配置（ds_config）并使用DeepSpeedEngine进行封装，而actor默认情况下将使用DeepSpeedHybridEngine进行封装，DeepSpeedHybridEngine的简单介绍可见下方
3. 最终得到1个携有所有相关模型的对象rlhf_engine

其对应的代码如下

1. ​step3_rlhf_finetuning中的main.py中，可以看到调用了DeepSpeedRLHFEngine

   # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
   """
   使用DeepSpeedRLHFEngine类直接初始化模型
   当然其内部仍旧调用了“create_hf_model”方法来读取模型，
   但其中实现了更为精细的DeepSpeed控制
   """
   rlhf_engine = DeepSpeedRLHFEngine(···)

2. 而DeepSpeedRLHFEngine的实现在step3_rlhf_finetuning/rlhf_engine.py中，涉及到actor、ref、critic、reward等4个模型的初始化

   # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py
   class DeepSpeedRLHFEngine():
       def __init__(···):
           """
           加载模型并进行DS封装
           1. actor与ref（以及actor_ema）通常都初始化自phase1训练所得的模型；
           2. critic与reward通常都初始化自phase2训练所得的模型
           根据它们的入参就能知道
           """
           ···
   
           """此处的actor是模型经过DeepSpeed封装后得到的DeepSpeedHybridEngine对象"""
           self.actor = self._init_actor(actor_model_name_or_path)
   
           """此处的reference是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
           self.ref = self._init_ref(actor_model_name_or_path)
   
           self.actor_ema = None
           """如果开启了ema，则初始化并封装ema"""
           if self.args.enable_ema:
               """此处的ema是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
               self.actor_ema = self._init_ema(actor_model_name_or_path)
   
           """此处的critic是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
           self.critic = self._init_critic(critic_model_name_or_path)
   
           """此处的reward是模型经过DeepSpeed封装后得到的DeepSpeedEngine对象"""
           self.reward = self._init_reward(critic_model_name_or_path)

3. 然后，其中actor的初始化细节是这样的

   # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py
   def _init_actor(self, actor_model_name_or_path):
       """
       初始化actor并使用DeepSpeedHybridEngine封装
       :param actor_model_name_or_path: phase1训练好的actor模型路径
       :return: 经DeepSpeedHybridEngine封装的actor
       """
   	···
   
       """
       DS Config
       根据传参构建ds config，
       与其他相关模型不同的地方在于，如果传参指定启用了enable_hybrid_engine，
       那么HybridEngine将作用于actor，对actor进行封装，
       因为HybridEngine可以使得模型可以在训练与推理两种模式中进行自动切换，
       同时享有训练与推理的优化，
       这对于既需要进行推理生成、又需要进行训练的actor来说是有增益作用的。
       """
       ds_config = get_train_ds_config(···,
           enable_hybrid_engine=self.args.enable_hybrid_engine,
           ···)
   	···
   
       # Model
       """使用CausalLM结构载入模型及权重，实例化actor"""
       actor_model = create_hf_model(
           model_class=AutoModelForCausalLM,
           model_name_or_path=actor_model_name_or_path,
           ds_config=ds_config,
           ···)
   
       # LoRA
       """如果开启LoRA训练则添加LoRA旁路"""
       if self.args.actor_lora_dim > 0:
           actor_model = convert_linear_layer_to_lora(···)
           if self.args.only_optimize_lora:
               actor_model = only_optimize_lora_parameters(actor_model)
   
       # Optimizer
       """实例化优化器：分组权重衰减等"""
       AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
       optim_params = get_optimizer_grouped_parameters(
           actor_model, self.args.actor_weight_decay)
       optim = AdamOptimizer(optim_params,
                             lr=self.args.actor_learning_rate,
                             betas=(0.9, 0.95))
   
       # LR Scheduler
       """实例化学习率调度器"""
       lr_scheduler = get_scheduler(
           name=self.args.lr_scheduler_type,
           optimizer=optim,
           num_warmup_steps=self.args.num_warmup_steps,
           num_training_steps=self.num_total_iters,
       )
   	
   	"""
       DeepSpeedEngine封装
       若ds_config中定义了启用HybridEngine，
       则返回的actor_engine不仅是个DeepSpeedEngine实例，
       确切地说还是个DeepSpeedHybridEngine实例，集成有HybridEngine的优化
       """
       actor_engine, *_ = deepspeed.initialize(model=actor_model,
                                               optimizer=optim,
                                               lr_scheduler=lr_scheduler,
                                               config=ds_config)
       ···
       return actor_engine

   其余ref、actor_ema、critic、reward的初始化几乎同理，只是ds_config设置不同，但最终都将返回经DeepSpeedEngine封装的对象

3.3.2 reward_score和values的区别及对经验数据的获取

3.3.2.0 经验数据的获取：step3_rlhf_finetuning/ppo_trainer.py

类似下图所示，DeepSpeed-Chat本阶段经验数据的获取过程如下：

1. 备有prompt数据（prompt_input_ids，prompt_attention_mask）；
2. 使用当前actor对prompt进行answer生成，得到完整对话序列seq（即上图中的sequence）；
3. 将seq输入至当前actor，输出得到当前(旧)策略logits（即上图中的action_logits），取对数logprobs；
4. 将seq输入至ref/SFT，输出得到baseline策略ref_logits（即上图中的sft_logits），取对数ref_logprobs；
5. 将seq输入至reward/RM，输出得到环境奖励reward_score（即上图中的r(x,y)）；
6. 将seq输入至当前critic，输出得到当前(旧)价值估计values（即上图中的value）；
7. 至此，用于进行PPO训练的各个基本经验数据已经获取齐全，至于图示的adv、reward等数据，在DeepSpeed-Chat中，于具体训练过程才开始计算
   (顺带要提的是，图示的reward指InstructGPT所提及的“KL Reward”：为了防止对phase2学习到的reward过度自信，引入了SFT与logits的KL散度作为惩罚的Reward，具体下文会细致阐述)

相关代码实现可见下方代码块

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def generate_experience(self, prompts, mask):
    """
    生成经验
    :param prompts: prompt input ids，(bs, max_prompt_len)
    :param mask: prompt attention mask, (bs, max_prompt_len)
    :return:
    """
    
    """将actor、reference、critic、reward转换为eval模式"""
    self.eval()
    
    """
    seq.shape: (seq_bs, max_seq_len)
    seq_bs指：排除较短answer后的batch_size。
    所谓“较短answer”在默认设定中是“序列长度小于1的answer”，
    短answer的seq都被滤掉了，
    所以可能batch_size会比之前小，
    但这个可能性极低，DS-C认为只有在使用未经phase1训练的模型来生成才会出现该情况。
    
    _generate_sequence()更具体的细节可见后续详解
    """
    seq = self._generate_sequence(prompts, mask)
    
    """将actor、critic转换为train模式，因为后续两者仍需要进行训练"""
    self.train()
    
	···
	
    with torch.no_grad():
    	"""
        经验采集：这部分其实就是在获取计算phase3损失函数所需的内容
        1. actor:（旧）策略-output.logits
        2. reference:SFT策略-output_ref.logits
        3. reward:奖励-reward_score，InsructGPT中的r_\theta
        4. critic:（旧）价值估计-values
        """
        output = self.actor_model(seq, attention_mask)
        output_ref = self.ref_model(seq, attention_mask)

        # (seq_bs, max_seq_len, vocab_size)
   		logits = output.logits

    	# (seq_bs, max_seq_len, vocab_size)
    	logits_ref = output_ref.logits
		
		"""价值函数的forward_value()更具体的细节下文马上讲解 """
		"""reward_score取的是answer最后一个token的value"""
        # reward_score.shape: (seq_bs,)
        reward_score = self.reward_model.forward_value(
            seq, attention_mask,prompt_length=self.prompt_length)['chosen_end_scores'].detach()

        """critic_model.forward_value(return_value_only=True)
将返回shape为(seq_bs, max_seq_len)的序列各token的value"""
        # 相当于就输出了旧价值values序列
        values = self.critic_model.forward_value(
        	seq, attention_mask, return_value_only=True).detach()[:, :-1]

    # 返回的dict是“进行PPO所需要使用的一组数据”
    # prompts.shape: (bs, max_prompt_len)
    # logits[:, :-1, :].shape: (seq_bs, max_seq_len - 1)
    # seq[:, 1:].shape: (seq_bs, max_seq_len - 1)
    # gather_log_probs()相当于输入logits和labels，对logits进行log_softmax后取出对应label位置的logit值
    # 因此logprobs.shape: (seq_bs, max_seq_len - 1)，ref_logprobs.shape: (seq_bs, max_seq_len - 1)
    # values.shape: (seq_bs, max_seq_len - 1)
    # rewards.shape: (seq_bs,)，reward_score在InstructGPT中就是r_\theta
    # input_ids.shape: (seq_bs, max_seq_len)
    # attention_mask.shape: (seq_bs, max_seq_len)
    """gather_log_probs()更具体的细节可见后续详解。"""
    return {
        'prompts': prompts,
        'logprobs': gather_log_probs(logits[:, :-1, :], seq[:, 1:]),
        'ref_logprobs': gather_log_probs(logits_ref[:, :-1, :], seq[:,
                                                                    1:]),
        'value': values,
        'rewards': reward_score,
        'input_ids': seq,
        "attention_mask": attention_mask

接下来，有三点需要重点说明

3.3.2.1 seq的生成：step3_rlhf_finetuning/ppo_trainer.py

对于本次batch的prompt，将输入至当前actor（对于即将根据经验数据迭代得到的actor来说，此时的“当前actor”可以认为是“旧策略网络”）来生成answer（如下图所示），然后将prompt与answer进行拼接得到seq

示例代码如下

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def _generate_sequence(self, prompts, mask):
    """
	生成seq
    """
	···
	"""
    获取prompt拼接上answer后的最大长度，实际上相当于max_seq_len，用于对生成长度做限制
    """
    max_min_length = self.max_answer_seq_len + prompts.shape[1]

    with torch.no_grad():
        """调用actor，输入input_ids和attention_mask进行生成"""
        seq = self.actor_model.module.generate(prompts,
                                               attention_mask=mask,
                                               max_length=max_min_length,
                                               min_length=max_min_length)

    """下方操作是为了过滤掉只有极短answer（有效长度小于1）的seq"""
    batch_size = seq.shape[0]

    """prompt长度：实际上就是max_prompt_len"""
    prompt_length = prompts.shape[1]

    """取出answer部分，此时还含有pad token"""
    ans = seq[:, prompt_length:]

    """统计answer的有效长度（去掉pad token后的长度）"""
    valid_ans_len = (ans != self.tokenizer.pad_token_id).sum(dim=-1)

    """排除较短（此处设置为有效长度小于1）的answer，余下的answer将被存入out_seq作为最终返回"""
    out_seq = []
    for i in range(batch_size):
        if valid_ans_len[
                i] <= 1:  # if the answer is shorter than 1 token, drop it
            continue
        else:
            out_seq.append(seq[i:i + 1])
    out_seq = torch.cat(out_seq, dim=0)  # concate output in the batch dim

    # out_seq.shape: (valid_batch_size, max_seq_len)
    return out_seq

3.3.2.2 奖励reward_score和价值估计values的区别：utils/model/reward_model.py

“奖励/环境奖励/reward_score”主要是为整个对话序列给出一个奖励值/做出评分
“价值估计/values”是为对话序列中的每一个位置都给出价值预测，是与时间步/状态紧密相关的

举个例子，有对话序列 seq=[11, 22, 33, 44, 55, 66, 0, 0, 0, 0]，其conversation_rewards = [2.01, 0.23, 2.89, 0.66, 0.33, 2.25, 0.36, 0.99, 1.32, 1.62]

• 其奖励reward_score只会是1个标量，具体是最后一个有效token所对应的reward_score_seq = 2.25；
• 其价值估计values是1维数组，如values_seq=[0.21, 1.26, 2.52, 0.03, 0.59, 1.55, 1.75, 2.12, 2.22, 1.32]

如下，奖励模型的模型类RewardModel中实现了用于取到环境奖励与价值估计的方法，即forward_value()，有两点需要重点强调下

1. 如果在本经验生成阶段3.3.2 节中调用这个forward_value的话，那得到的values是旧的

           # 相当于就输出了旧价值values序列
           values = self.critic_model.forward_value(
           	seq, attention_mask, return_value_only=True).detach()[:, :-1]

   如果在下文「3.3.4.4 最后计算价值损失」节中计算loss的过程中，调用这个forward_value的话，得到的values就是新的

       # 且此时因为是计算价值损失，所以这里计算的是新价值估计
       value = self.critic_model.forward_value(**batch,
                                               return_value_only=True,
                                               use_cache=False)[:, :-1]

2. 其与该类RewardModel训练所用的的另一个方法forward()不同，另一个方法forward() 主要实现了环境奖励的获取与排序损失的计算
   总之，RewardModel这个类既实现了训练时用的forward() 的方法，也实现了推理时用的forward_value() 方法

最终，forward_value在RewardModel这个类中 是如下具体实现的

# applications/DeepSpeed-Chat/training/utils/model/reward_model.py
class RewardModel(nn.Module):

    def __init__(self, base_model, tokenizer, num_padding_at_beginning=0):
        ···
    ···
    def forward(···):
    	"""forward()在之前“2.3.3 整个对话的reward设计和成对排序损失”中已经进行过详解，且与此处所述内容无关，此处不再赘述"""
        ···

    def forward_value(···, return_value_only=False, ···):
        """
        和forward的差别在于：forward需要针对输入的chosen-rejected对计算排序损失并返回
        而forward_value只需要考虑一个输入，然后返回分值
        说白了，forward的输入是数据对，因为要计算数据对的排序损失，而forward value的输入是单个数据，直接推理出其分值
        至于参数return_value_only: 如果设置为True，则在计算出values（在序列上每个位置的分值预测）后直接返回
        """
        
        """经过主干网络正向传播得到输出"""
        transformer_outputs = self.rwtranrsformer(···)

        # hidden_states.shape: (bs, max_seq_len, hidden_size)
        hidden_states = transformer_outputs[0]

        """将隐状态特征传入线性层v_head输出得到分值"""
        # values.shape: (bs, max_seq_len)
        values = self.v_head(hidden_states).squeeze(-1)
        
        if return_value_only:
        	"""
			如果传参中预设了“return_value_only=True”，
			那么将直接返回 values: (bs, max_seq_len)
			"""
            return values
        else:
        	"""否则还将进一步取得reward_score"""
            # 相当于为true  返回values序列，为false 返回values序列和reward标量值 
            bs = values.size(0)
            seq_len = input_ids.shape[1]
            chosen_end_scores = []
            for i in range(bs):
            	···
                # value.shape: (max_seq_len,)
                value = values[i]

                """c_ind即为prompt之后的序列片段中，第一个pad_token的index"""
                c_ind = ···

                """取c_ind的前一个index（实际上就是answer的最终位置）作为reward_score"""
                ···
                chosen_end_scores.append(value[c_ind - 1])
            
            """返回values和reward_score"""
            return {
                "values": values,
                "chosen_end_scores": torch.stack(chosen_end_scores),
            }

3.3.2.3 策略模型logits的进一步处理

策略模型（actor、ref/SFT）所输出logits的shape为(bs, max_seq_len, vocab_size)，然而计算KL散度惩罚、重要性权重时并不需要对所有vocab的logits进行计算，仅需要对groundtruth项（seq各个token对应的项）的logits进行计算即可

batch_size = 1
max_seq_len = 4
vocab_size  = 3

logits = [
          [[1.23, 2.11, -0.56], 
           [-1.52, -1.11, 1.66], 
           [0.32, 0.13, 1.55], 
           [-0.55, -0.23, -1.62]]
         ]

seq = [
       [2, 2, 0, 1]
      ]

对于CausalLM来说，logits第t个时间步的置信值是为了预测第t+1步的seq token，因此logits[, :-1, :]与seq[:, 1:]才是“预测与标签”的关系：

logits[, :-1, :] = [
                      [[1.23, 2.11, -0.56], 
                       [-1.52, -1.11, 1.66], 
                        [0.32, 0.13, 1.55]]
                    ]
seq[:, 1:] = [
              [2, 0, 1]
             ]

只需要从预测中根据对应标签取出logits即可，以上述例子为例，最终取出的结果probs为

probs = [
             [-0.56, -1.52, 0.13]
            ]

因此DeepSpeed-Chat定义了函数gather_log_probs()来对输出的logits进行后处理，以获取对数化后的结果log_probs

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def gather_log_probs(logits, labels):
    """
    相当于输入logits和labels，对logits进行log_softmax后取出对应label位置耳朵logit值
    :param logits: (bs, seq_len, vocab_size)
    :param labels: (bs, seq_len)
    :return: log_probs_labels.squeeze(-1): (bs, seq_len)
    """

    # log_probs.shape: (bs, seq_len, vocab_size)
    log_probs = F.log_softmax(logits, dim=-1)

    """
    此处gather()可以根据labels（index）来从log_probs中获取对应index的值
    总的来说就是取出logits中对应labels数值位置的值
    log_probs_labels.shape: (bs, seq_len, 1)
    """
    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
    return log_probs_labels.squeeze(-1)

3.3.3 PPO训练数据管理-MiniDataset：utils/data/data_utils.py

最开始的时候载入过一次Dataset，但刚开始载入的Dataset针对的是全部训练数据的管理，而此时使用的MiniDataset主要针对PPO训练迭代所使用的数据进行管理。PPO训练前的数据管理流程可以理解为：

1. 首先，Dataloader从Dataset中取出：1个prompt_batch的无监督数据，和1个prompt_batch的prompt数据
   (注：这里的无监督数据就是为了实现那个ptx项，至于为何要有这个ptx项，原因在于：无监督训练使得模型具备基本的生成通顺语句的能力，而在RLHF阶段引入ptx 可以让模型在追求人类偏好的同时也不要遗忘基本生成能力)
   对于后者，如果使用1个prompt_batch的prompt数据进行经验采集，将得到1个prompt_batch的经验数据
2. 之后，1个prompt_batch的无监督数据、1个prompt_batch的经验数据将被送入各自的MiniDataset实例进行管理：1个prompt_batch将被分成数个ppo_batch，供PPO训练进行数次迭代，如下代码所示(来自step3_rlhf_finetuning/main.py)

   # applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
   """经验数据以及无监督数据都将被MiniDataset所管理"""
   exp_mini_dataset = MiniDataset(···)
   unsup_mini_dataset = MiniDataset(···)  
   
   # out为经验数据     
   out = trainer.generate_experience(···)
   exp_dataset = exp_mini_dataset.add(out)
   unsup_dataset = unsup_mini_dataset.add(batch_unsupervised)

上述第2步就是MiniDataset所要做的事，而MiniDataset这个类被定义在utils/data/data_utils.py中，分别执行了以下三个操作：

1. seperate()：细分为ppo_batch数据，其具体实现的代码为

   # applications/DeepSpeed-Chat/training/utils/data/data_utils.py
   class MiniDataset:
       def __init__(self, max_size, small_batch_size):
           """
           :param max_size: batch数。通常此处指“用于给actor做生成的prompt的batch数（注意是batch数不是batch_size）”
           :param small_batch_size: batch size。通常此处指“PPO训练的batch_size”。
           """
           self.dataset = []
           self.max_size = max_size
           self.small_batch_size = small_batch_size
   
       def seperate(self):
       	"""维护1个small_dataset"""
           small_dataset = []
   
           # 从self.dataset中逐个取batch
           for large_batch in self.dataset:
               """判断batch的数据类型（列表/元组/字典），
               根据数据类型取其batch_size，赋值给large_size"""
               if type(large_batch) == list or type(large_batch) == tuple:
                   large_size = len(large_batch[0])
               elif type(large_batch) == dict:
                   large_size = len(large_batch[list(large_batch.keys())[0]])
               else:
                   large_size = len(large_batch)
               """
   
               以下部分代码略微抽象，需要举例说明
               - 比如prompt的batch_size设置为3，PPO训练用的batch_size设置为4，则最后能取来用、存入small_dataset的也就只有3条数据
               - (因为生成用的dataloader只采样出了3条，最多也就只有3条)
   
               - 比如prompt的batch_size设置为5，PPO训练用的batch_size设置为4，则最后能取来用、存入small_dataset的就是2组数据
               - (第1组为idx0,idx1,idx2,idx3共4条数据、第2组为idx4共1条数据)
   
               - 比如prompt的batch_size设置为9，PPO训练用的batch_size设置为4，则最后能取来用、存入small_dataset的就是3组数据
               - ([0,1,2,3],[4,5,6,7],[8]）
               """
               for i in range(0, large_size, self.small_batch_size):
                   if type(large_batch) == list or type(large_batch) == tuple:
                       small_dataset.append(
                           [x[i:i + self.small_batch_size] for x in large_batch])
                   elif type(large_batch) == dict:
                       small_dataset.append({
                           k: v[i:i + self.small_batch_size]
                           for k, v in large_batch.items()
                       })
                   else:
                       small_dataset.append(large_batch[i:i + self.small_batch_size])
           """清空self.dataset"""
           self.free()
           
           """返回最终取用的数据，该ppo_batch数据将用于ppo训练迭代"""
           return small_dataset

2. add()：获取batch (prompt_batch)数据；

       def add(self, data):
           """
   		在最开始的时候可以传参预设“生成X个batch再进行PPO训练”，
   		此处的max_size就是其中的X，
   		如果少于max_size则将batch数据加入至MiniDataset中，
   		直至达到max_size个batch
   		"""
           if len(self.dataset) < self.max_size:
               self.dataset.append(data)
               if len(self.dataset) == self.max_size:
                   """
                   seperate()主要实现了
                   1. 在batch的基础上，再细分ppo_batch并返回
                   2. 清空MiniDataset中的数据
                   """
                   return self.seperate()
               else:
                   return None
           else:
               raise ValueError(
                   "The dataset is full but we did not stop it. There is a bug in the code."
               )

3. free()：清空获取到的batch数据并返回ppo_batch数据

       def free(self):
           """清空self.dataset中的数据"""
           self.dataset = []

3.3.4 AC架构下的PPO训练：在加了β惩罚且截断后的RM之下，通过经验数据不断迭代策略且估计value

对于采集到的一批经验数据，使用MiniDataset处理成多批ppo_batch数据，供相关模型进行多次训练迭代

DeepSpeed-Chat中所设置的ppo_epochs，从强化学习的角度来说，实际上代表的是一批经验数据的复用次数：

• 假如ppo_epochs设置为1，训练时，引入的这批经验数据在经过1次全遍历后，将被直接弃置，随之进行下一轮prompt_epoch，届时将重新采集新的一批经验数据
• 假如ppo_epochs设置为n，训练时，引入的这批经验数据将被遍历n次才被弃置，即相当于这批经验数据被复用了n次用于off-policy训练

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py，以下是其中的第470-490行
for ppo_ep in range(args.ppo_epochs):
    """ppo_epoch循环"""
    for i, (exp_data, unsup_data) in enumerate(zip(exp_dataset, unsup_dataset)):
        """
        ppo_step循环：
        从MiniDataset返回的数据中，
        取1个ppo_batch的经验数据和无监督数据来训练
        """

        """经验数据训练，返回actor_loss和critic_loss"""
        actor_loss, critic_loss = trainer.train_rlhf(exp_data)

        """累加本ppo_step的指标，后续将除以内层迭代次数计算均值"""
        actor_loss_sum += actor_loss.item()
        critic_loss_sum += critic_loss.item()
        average_reward += exp_data["rewards"].mean()

        """无监督数据训练"""
        if unsupervised_training_enabled:
            """返回无监督损失"""
            unsup_loss = trainer.train_unsupervised(unsup_data, 
            											args.unsup_coef)
            """累加本ppo_step的无监督损失，后续将除以内层迭代次数计算均值"""
            unsup_loss_sum += unsup_loss.item()

        """PPO训练迭代次数（ppo_step）+1"""
        inner_iter += 1

        """是否启用指数移动平均技术"""
        if args.enable_ema:
            moving_average(rlhf_engine.actor,
                           rlhf_engine.actor_ema,
                           zero_stage=args.actor_zero_stage)

	"""打乱数据供off-policy复用"""
    random.shuffle(exp_dataset)
    random.shuffle(unsup_dataset)

1次PPO训练由train_rlhf()方法进行管理，其内部主要实现了「注，以下的内容如果有不太理解的，可以结合《ChatGPT技术原理解析：从RL之PPO算法、RLHF到GPT4、instructGPT》一文中的3.2节增进理解」：

1. KL散度惩罚奖励old_rewards的计算，为了防止对phase2学习到的环境奖励过度自信，所以加入了KL散度惩罚项：

   

2. 优势advantages和回报returns的计算
   其中包括本框架在内多数框架的advantages实现并非纯粹使用TD-error，而是在TD-error的基础上结合了MC方法，也即GAE（广义优势估计）；
   对于全长为的轨迹来说，其某个时间步的优势为（时，advantage完全使用MC方法；时，advantage完全使用TD-error方法）：
   
   至于回报returns就是奖励reward的累计，对于全长为的轨迹来说，其到达某个时间步 时的回报为

   

3. 在1个ppo_batch中，actor的损失计算公式为：
   
   其中，指的仅是“answer”部分的内容，并不包括prompt部分
4. 在1个ppo_batch中，critic的损失计算公式为：
   裁剪新价值估计，使其不至于太偏离采集经验时的旧价值估计，使得经验回放仍能有效：

   

   critic将拟合回报R：

   

   其中，指的仅是“answer”部分的内容，并不包括prompt部分，相当于强调的是“这个损失公式只计算answer部分，prompt部分的损失是不计入这个公式的”

  接下来，我们看下代码实现。为保证阅读的流畅性，七月在线ChatGPT课学员春天对其中的部分代码进行了调整，使得相应的函数代码衔接在其调用后方，便于具体对照其传参，从而辨析传入的新旧策略、新旧价值估计等
        且为更加一目了然，我又把代码拆分出来了几段，且加了一系列公式、图示、解释、说明，最终好结合“代码与图示”做更直观的分析，给你一个独一无二的通透解释!

3.3.4.1 首先是一系列定义以及对阶段二的reward加个KL惩罚

对阶段二的reward加个KL惩罚所对应的公式展开后为(来源于本博客内另一篇文章ChatGPT技术原理解析中的3.1.3 InstructGPT训练阶段3：如何通过PPO算法进一步优化模型的策略)

对应的图示为
『且有两点值得特别注意的是

1. 实际代码实现时，对RM做带β 的KL惩罚时，分子取的是为经验数据中的旧策略(如上公式所示，对应的π(RL')，当然即便分子是经验数据中的旧策略，β对应的惩罚比值依旧是新旧之比：π(RL')/π(SFT)，因为π(RL')虽一开始被π(SFT)初始化，但一步或多步之后π(RL') 即更新了，至于到底是一步还是多步之后更新，取决于上文中说过的ppo_epochs等于1还是n，至于代码实现中这个惩罚比值则对应的：旧策略/SFT策略 = log_probs/ref_log_probs)，非下图图1中有歧义的action_logits
2. 此外，实际代码实现时，既对RM加了KL惩罚，同时出于某种安全考虑的需要，又对RM做了截断reward_clip，但下图图1没把这个reward_clip体现出来』

对应的代码为

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def train_rlhf(self, inputs):
    """
    使用1个ppo_batch的经验数据，执行1次rlhf训练迭代
    """
    # prompt input ids
    prompts = inputs['prompts']

    # (旧)策略
    log_probs = inputs['logprobs']

    # SFT策略
    ref_log_probs = inputs['ref_logprobs']

    # RM奖励
    reward_score = inputs['rewards']

    # (旧)价值估计
    values = inputs['value']
    attention_mask = inputs['attention_mask']

    # seq input ids
    seq = inputs['input_ids']
	
	"""
    获取prompts的最后1个位置作为start
    比如prompt_len为256，start则为 256-1=255
    这个start主要是用于取出经验数据中的“非prompt”部分（也即“answer+padding”部分）
    """
    start = prompts.size()[-1] - 1
	
	"""
	action_mask相当于取 attention_mask除了第0个序列位置外的部分，
	需要注意的是：
	1. 多数情况下，包括此处在内的transformers风格代码中，
	attention_mask指的实际上是“padding_mask”而非“sequence_mask”；
	2. 之所以要进行[:, 1:]切片，是为了去除第0个位置从而与seq对齐，
	因此 action_mask.shape: (bs, max_seq_len - 1)
	3. 后续将被用于过滤掉pad token位置的信息
	4. 但实际上在后续的使用中，
	基本都会结合上方定义的start，从action_mask中再切片出“非prompt”部分，
	例如 action_mask[start:]，实际上就相当于取“非prompt”部分，
	action_mask[start:].shape: (bs, max_answer_len)
	"""
	action_mask = attention_mask[:, 1:]    
    ···

    """经验数据中的价值估计为“旧”价值估计"""
    old_values = values
    with torch.no_grad():
    	
    	###计算KL惩罚修正的奖励################################################
        """
        通过KL散度惩罚，以及r_\theta（来自phase-2的reward model）计算得到修正的奖励，
        注意此处的入参：
        1. log_probs为经验数据中的旧策略『特别注意这里 是针对旧策略log_probs(如上公式所示)，非图示中的action_logits』
        2. ref_log_probs为经验数据中的SFT策略
        3. reward_score为经验数据中的RM赋分
        """
        old_rewards = self.compute_rewards(prompts, log_probs,
                                           ref_log_probs, reward_score,
                                           action_mask)
        def compute_rewards(self, prompts, log_probs, ref_log_probs, reward_score,
                    action_mask):
		    """
		    计算实际rewards，涉及(旧)策略与SFT的KL散度惩罚、RM的reward
		    """
		    """计算经验采样时actor与SFT的KL散度惩罚"""
		    kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)
		    rewards = kl_divergence_estimate

		    """
		    找到answer的起始start：即prompt的最后1个token位置
		    比如prompts长度为256，answer的起始则为256-1=255
		    """
		    start = prompts.shape[1] - 1

			"""
			ends为batch中各个数据的最后1个有效token的index，
			每个数据的最末有效token位置很大可能是不一样的，
			因此ends是个数组
			"""
		    ends = ···

		    """
		    将RM得到的奖励值限定在一定范围，默认为(-5,5)
            相当于既对RM加了修正，同时又对RM做了截断reward_clip
		    """
		    reward_clip = torch.clamp(reward_score, -self.clip_reward_value,
		                              self.clip_reward_value)
			···

		    """
		    因为batch中每个数据的最末有效token位置很可能不一样，
		    所以无法通过矩阵来并行，需要使用for循环逐个数据处理
		    """
		    for j in range(batch_size):
		    	"""
		        KL_reward = KL + reward
		        加和只在最末有效token上进行
		        """
		        rewards[j, start:ends[j]][-1] += reward_clip[j]	
		    """返回KL rewards"""
		    return rewards

3.3.4.2 接下来，计算优势与回报：step3_rlhf_finetuning/ppo_trainer.py之get_advantages_and_returns

(再提一嘴，如本3.1节开头所述，其实因为优势adv以及KL惩罚奖励完全可以由基本经验数据计算得到，故adv和KL惩罚奖励在生成经验的阶段一步到位计算即可)

        ###计算优势与回报################################################
        """
        计算优势advantages和回报returns
		注意此处的入参：
        4. old_value为经验数据中的(旧)价值估计
        5. old_rewards为刚才计算得到的KL_reward
		"""
        advantages, returns = self.get_advantages_and_returns(
            old_values, old_rewards, start)
		def get_advantages_and_returns(self, values, rewards, start):
		    """
		    计算优势与回报
		    实现基本与上述公式相同
		    """
		    lastgaelam = 0
		    advantages_reversed = []
		    length = rewards.size()[-1]
		    
		    """反向遍历计算各个时间步的优势advantage"""
		    for t in reversed(range(start, length)):
		    	"""获取下个时间步的价值估计V_{old}(s_{t+1})"""
		        nextvalues = values[:, t + 1] if t < length - 1 else 0.0

		        """计算单步TD-error"""
                # gamma代表折扣因子γ，δ 表示TD误差
		        delta = rewards[:, t] + self.gamma * nextvalues - values[:, t]

		        """累计优势"""
		        lastgaelam = delta + self.gamma * self.lam * lastgaelam

		        """存储各个时间步的优势"""
		        advantages_reversed.append(lastgaelam)

		    """对逆序的优势列表进行正序处理，得到正常时间步排列的优势"""
		    advantages = torch.stack(advantages_reversed[::-1], dim=1)

		    """
		    return_t = adv_t + v_(s_t)
		    由优势计算得到回报
		    """
		    returns = advantages + values[:, start:]

		    """返回优势与回报"""
		    return advantages.detach(), returns

————————————————
另，我在 《ChatGPT技术原理解析：从RL之PPO算法、RLHF到GPT4、instructGPT》一文的3.2节中，举了一个完整计算优势值和回报值的例子，通过例子可以对整个计算过程更加一目了然
“为进一步帮助大家理清相关概念的每一个细节，我们来具体看下DSC到底是怎么计算returns的

在DSC中，优势函数被定义为对未来所有时间步的TD误差（Temporal-Difference Error，也就是这里的δ）进行折扣求和，这种计算优势函数的方法，被称为Generalized Advantage Estimation (GAE)

首先，我们需要计算TD误差序列，在这里，我们使用提供的V_old值序列和奖励序列进行计算：
values = [-0.2761, -2.3945, 0.1729, -0.0919, -0.0867, -0.0818, -0.0758]
KL_rewards = [-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -8.5999e-02]
γ=0.9
...”
且这个例子对应的计算代码为

# -*- coding: utf-8 -*-
"""
Created on Wed Jan 19 20:42:44 2022
@author: Remixa
"""
import torch

def get_advantages_and_returns(values, rewards, start, length, gamma, lam):
    lastgaelam = 0
    advantages_reversed = []
    length = rewards.size(-1)
    print("#" * 10 + "init:" + "#" * 10)
    print(f"start: {start}, length: {length},\ngamma: {gamma}, lam: {lam}\nvalues: {values}\nrewards: {rewards}")
    for t in reversed(range(start, length)):
        nextvalues = values[:, t + 1] if t < length - 1 else torch.zeros_like(values[:, t])
        delta = rewards[:, t] + gamma * nextvalues - values[:, t]
        lastgaelam = delta + gamma * lam * lastgaelam
        advantages_reversed.append(lastgaelam)
        print("#" * 10 + f"t = {t}" + "#" * 10)
        print(f"t = {t}:\nnextvalues: {nextvalues}\ndelta: {delta}\nlastgaelam: {lastgaelam}\nadvantages_reversed: {advantages_reversed}")
    advantages = torch.stack(advantages_reversed[::-1], dim=1)
    returns = advantages + values[:, start:]
    print("#" * 10 + "result:" + "#" * 10)
    print(f"advantages: {advantages}\nreturns: {returns}")
    return advantages.detach(), returns

if __name__ == '__main__':
    old_values = torch.tensor([[-0.2761, -2.3945, 0.1729, -0.0919, -0.0867, -0.0818, -0.0758]])
    old_rewards = torch.tensor([[-4.6873e-04, -3.1257e-04, 5.8591e-05, -5.5084e-03, -4.0741e-03, -5.5275e-03, -8.5999e-02]])
    start = 3
    length = 7
    gamma = 0.9
    lam = 0.95
    advantages, returns = get_advantages_and_returns(old_values, old_rewards, start, length, gamma, lam)

3.3.4.3 然后计算策略损失：step3_rlhf_finetuning/ppo_trainer.py之actor_loss_fn

对应的图示为

但看代码之前，有两点得提一下：

1. 新旧策略的重要性比值：logprobs/old_logprobs 体现出来了，但对新策略logprobs做的截断在上图中没有体现出来，毕竟我们已经知道，需要通过截断或KL散度约束这个比值，防止新策略相比旧策略更新过大
2. log_probs 为之前定义的经验数据中的(旧)策略，与之对应，actor的新策略则定义为 actor_log_probs，代码实现中的这4个定义的对应关系如下图所示(之所以如此，可能如春天所说，写策略损失函数的是一个人，写经验数据的又是另一个，总之 大家注意下，避免混淆绕晕)

对应的代码为

    ###计算actor损失并更新################################################
    batch = {'input_ids': seq, "attention_mask": attention_mask}

    """将seq经验数据输入至actor，进行自回归预测"""
    actor_prob = self.actor_model(**batch, use_cache=False).logits

    """取出probs，此处为新策略"""
    actor_log_prob = gather_log_probs(actor_prob[:, :-1, :], seq[:, 1:])

    """
    计算actor损失
    注意此处的入参：
    1. actor_log_probs为方才刚输出的新策略
    2. log_probs为之前定义的经验数据中的(旧)策略
    3. advantages为之前计算出的优势
    """
    actor_loss = self.actor_loss_fn(actor_log_prob[:, start:],
                                    log_probs[:, start:], advantages,
                                    action_mask[:, start:])
	def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask):
	    """计算actor的损失"""
	    
	    """
	    重要性采样权重计算：ratio = exp(log(new)-log(old)) 
	    """
	    log_ratio = (logprobs - old_logprobs) * mask
	    ratio = torch.exp(log_ratio)

		"""计算策略梯度损失的2个情况：加权优势 与 裁剪加权优势"""
	    pg_loss1 = -advantages * ratio
	    pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange,
	                                         1.0 + self.cliprange)

	    """
		从策2种情况中选择损失较大者作为真正的损失，
		并且基于ppo_batch内所有数据的所有有效时间步计算平均损失值
		"""
	    pg_loss = torch.sum(torch.max(pg_loss1, pg_loss2) * mask) / mask.sum()
	    return pg_loss

	"""actor反向传播、更新参数"""
    self.actor_model.backward(actor_loss)
    self.actor_model.step()

3.3.4.4 最后计算价值损失：step3_rlhf_finetuning/ppo_trainer.py之critic_loss_fn

    ###计算critic损失并更新################################################
    """将seq经验数据输入至critic，预测得到新价值估计"""
    # 调用的forward_value即是上文「9.3.2.2 奖励reward_score和价值估计values的区别」中分析的那个
    # 且此时因为是计算价值损失，所以这里计算的是新价值估计
    value = self.critic_model.forward_value(**batch,
                                            return_value_only=True,
                                            use_cache=False)[:, :-1]

    """
    计算critic损失
    注意此处的入参：
    1. values为方才刚输出的新价值估计
    2. old_values为经验数据中的(旧)价值估计
    3. returns为之前计算出的回报
    """
	critic_loss = self.critic_loss_fn(value[:, start:], old_values[:,start:],
                                      returns, action_mask[:, start:])
	def critic_loss_fn(self, values, old_values, returns, mask):
	    """计算价值损失"""

	    """裁剪当前新values，使得其不至于太偏离经验采样阶段的旧values"""
	    values_clipped = torch.clamp(
	        values,
	        old_values - self.cliprange_value,
	        old_values + self.cliprange_value,)
	    
	    """计算当前values与回报的L2 Loss"""
	    vf_loss1 = (values - returns)**2

	    """计算裁剪后的当前values与回报的L2 Loss"""
	    vf_loss2 = (values_clipped - returns)**2

	    """
	    选择损失较大者作为真正的损失，
		并且基于ppo_batch内所有数据的所有有效时间步计算平均损失值，
		此外critic损失项的系数为0.5。
	    """
	    vf_loss = 0.5 * torch.sum(
	        torch.max(vf_loss1, vf_loss2) * mask) / mask.sum()
	    return vf_loss
	    
    """critic反向传播、更新参数"""
    self.critic_model.backward(critic_loss)
    self.critic_model.step()
	
	
	"""本次ppo_step将返回actor_loss和critic_loss供指标统计"""
    return actor_loss, critic_loss

顺带说下，在进行phase3的RLHF训练时，为使得模型在学习人类偏好的过程中仍能保有预训练模型解决任务的性能，引入了传统的自回归语言建模进行联合训练

对应的示例代码为

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py
unsup_loss = trainer.train_unsupervised(unsup_data, args.unsup_coef)
def train_unsupervised(self, inputs, unsup_coef):
    """
    1个ppo_batch的无监督训练
    :param inputs: dict：input_ids, attention_mask, labels
    :param unsup_coef: 无监督损失系数
    """
    """确保actor处于训练模式，否则将返回报错"""
    self._validate_training_mode()

    """actor进行常规的CausalLM训练"""
    outputs = self.actor_model(**inputs, use_cache=False)
    loss = outputs.loss
    """反向传播、更新参数"""
    self.actor_model.backward(unsup_coef * loss)
    self.actor_model.step()

    return loss

---

最后，再次引用学员春天的几点总结：

1. “RLHF的训练涉及到强化学习，训练过程对超参数的设置极其敏感，DeepSpeed-Chat团队在尝试了多种参数设置后，最终默认设置了per_device_train_batch_size(即prompt_batch_size) = per_device_mini_batch_size(即ppo_batch_size)，且生成1个prompt_batch就立刻开始训练——这样一来，实际上在进行的就是On-Policy强化学习，采集一次、学习一次，数据利用率并不高"
2. 此外，DeepSpeed-Chat团队还发现为无监督训练的损失设置系数（unsup_coef）也非常困难，训练过程会变得更加震荡，不过团队也没有花费太多精力在调整这个系数参数上

当然这些都并不是最佳的超参数配置，DeepSpeed-Chat团队仍鼓励用户多做尝试并分享出自己的调参经验