前言

  本篇为上中下三篇章的【下篇】，接续自【中篇】。主要针对整个DeepSpeed-Chat框架中最为复杂的第三阶段进行详解，其中涉及到部分InstructGPT所述相关原理的实践，基于其代码将更易于理解诸多原理中一笔带过的话题，如“用于经验采集的策略网络到底是SFT还是Actor”“Critic的迭代又是如何实现的”等等。
  尽管这是讨论DeepSpeed-Chat技术为主的文章，但还是不得不借用ColossalChat（另一个实现了RLHF Pipeline的开源项目，项目地址）绘制的流程示意图（下图），因为它所描绘的第三阶段训练流程，非常详细且与DeepSpeed-Chat的实现过程基本一致，而DeepSpeed-Chat本身给出的示意图实在太过简略（见【中篇】头图）。

  相信结合这张示意图来学习phase3效果更佳。

3 phase-3: RLHF Finetuning

3.1 训练数据样例

3.1 基本数据

数据格式名称	说明	样例
prompt	对当前情境的描述，为模型生成提供指令输入信息，可以理解为通俗含义上的“问句”，适用于phase3。	"Human: Please tell me about Microsoft in a few sentence? Assistant: "

3.2 经验数据

数据格式名称	说明	样例
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	reward/RM基于seq输出的对整个对话的（环境）奖励。	shape: (seq_bs,)
attention_mask	用于滤掉非有效元素。	shape: (seq_bs, max_seq_len)

  各个框架对于经验数据的定义不完全相同，例如ColossalChat定义的经验数据还比此处多了项“adv”和“reward”（此reward非彼reward，ColossalChat的reward指的是“经过KL散度修正后的KL_Reward”），但本质上都是同理的，只是框定的范围不同，因为adv（优势函数Adventage）和KL_Reward完全可以由已有项logprobs、ref_logprobs、reward、value计算得到。

  从代码效率的角度来考量，ColossalChat的经验数据定义相对更严谨些，因为优势以及KL惩罚奖励完全可以由基本经验数据计算得到，在生成经验的阶段一步到位计算即可；而DeepSpeed-Chat中将其安排在训练阶段来计算，每次PPO迭代才计算，优势和KL惩罚奖励是基于基本经验数据计算得到的，而基本经验数据在生成经验阶段已经确定了，所以即使是在不同的PPO迭代中，优势和KL惩罚奖励也是不变的，因此DeepSpeed-Chat对adv以及KL惩罚奖励进行了重复计算，这个环节的计算顺序后续（编辑日期2023.05.19）相关团队应该会做出调整。

3.2 训练过程

在此简单讲述UML时序图的元素含义：
- 箭头表示信息传递：实线表示调用，虚线表示返回；
- alt表示假设分支，其后方“[]”中的内容表示“条件”；
- loop表示循环；
- 淡蓝色区域即为高亮部分。

phase3的大致训练过程如UML时序图所示（“括号序号”与UML时序图的“圈序号”对应）：

1. 载入tokenizer(1-2)；
2. 获取Dataset并实例化DataCollator(3-9)：获取用于采集经验的prompt的Dataset(4-5)，如果启用了无监督训练，则再获取无监督数据的Dataset(6-7)，并且实例化DataCollator用于进一步对加载的数据进行整理；
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）:
   1. 大轮次step（prompt_step）：
      1. 使用MiniDataset.add()分配PPO训练轮次所用的训练数据(17-22)：
         1. 无监督数据部分直接使用MiniDataset进行分配(17-18)；
         2. 经验数据部分将由各相关模型采集得到(19-20)，然后再使用MiniDataset对经验数据进行分配(21-22)；
      2. 小轮次epoch（ppo_epoch）：
         1. 小轮次step（ppo_step）：
            1. 训练得到actor和critic的loss（过程中已经进行过参数更新）(23-24)，如果启用了无监督训练，则也进行无监督训练并返回loss(25-26)；
         2. 当启用指数移动平均技术（ema）时，进行ema(27)；
      3. 保存模型(28)。

3.3 关键代码详解

上述过程存在几个值得关注的地方（即文字描述加粗、UML时序图高亮的部分）：

• 数据集及Datacollator，知悉无监督数据以及prompt数据在读取后、输入模型前的形式；
• 各相关模型初始化细节；
• 经验数据的获取；
• PPO训练数据管理；
• PPO训练过程。

以下将对相关部分的源码进行讲解。

3.3.1 读取数据集及Datacollator

3.3.1.1 prompt数据集的读取

  至于数据集读取，prompt数据集将如【上篇】所述从本地进行读取载入，在原先的缓存预处理中，使用非padding的tokenizer对prompt进行处理后，还使用了flip操作将得到的input_ids以及attention_mask进行了翻转倒序，并且在当时也提前解释了原因：主要是便于进行前侧padding的操作。后续经过data_collator“倒序的数据经过padding后再flip回来”的操作，pad_token将会位于前侧，而进行生成采样的actor将能接续prompt的内容进行自回归生成（更具体可参考【上篇】的0.3 板块问题 第1点）。

3.3.1.2 DataCollatorRLHF

  正如上方所述，phase3使用的data_collator实例化自DataCollatorRLHF，该类主要实现了“padding至max_prompt_len（默认为max_seq_len的一半），然后进行flip”。

3.3.1.3 无监督数据集的读取

  无监督数据集主要是进行了分块处理，将无监督语料全部拼接起来得到一个极长的序列，使用max_seq_len大小的滑窗对长序列进行分块，每个分块将作为1条无监督数据。

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py 
def create_datasets(···, train_phase=3):"""获取Dataset和实例化dataloader"""···"""返回读取到的prompt数据，该数据为经由tokenizer处理（tokenize但未padding），且本地存储后的 input_ids 和 attention_mask 数据。并且在【上篇】有所提及：phase3存储的数据是经过flip翻转的、是倒序的，后续将在data_collator中先padding后再flip回正序，这将使得pad_token位于数据前侧。"""prompt_train_dataset, _ = create_prompt_dataset(···, train_phase, ···)if unsupervised_training_enabled:"""如果启用无监督训练，则获取无监督数据，并将其处理成分块形式，每块为1条数据，为max_seq_len长度"""unsupervised_train_dataset = get_unsupervised_data(args, tokenizer)else:unsupervised_train_dataset = None"""实例化数据整理器data_collator"""data_collator = DataCollatorRLHF(···)"""实例化数据加载器dataloader并且使用data_collator整理读取到的prompt数据（如上述所说：先padding后flip）"""prompt_train_dataloader = DataLoader(···, collate_fn=data_collator, ···)if unsupervised_training_enabled:"""如果启用无监督训练，则实例化无监督数据加载器"""unsupervised_train_dataloader = DataLoader(···)else:"""如果未启用无监督训练，也仍实例化一个空的数据加载器，因为多数后续代码入参接口都支持同时输入prompt数据与无监督数据，这一步是基于后续传参的安全性考虑"""unsupervised_train_dataloader = [None] * len(prompt_train_dataloader)

3.3.2 初始化各相关模型

3.3.2.1 模型初始化过程

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

1. 读取模型，虽然同样支持直接从huggingface hub拉取相应模型，但通常都是从本地路径读取phase1、phase2训练好的模型：
   1. actor、ref/SFT以及actor_ema（如果开启了ema）通常都初始化自phase1训练所得的模型；
   2. critic与reward通常都初始化自phase2训练所得的模型。
2. 为各个相关模型设定不同的DeepSpeed配置（ds_config）并使用DeepSpeedEngine进行封装，而actor默认情况下将使用DeepSpeedHybridEngine进行封装，DeepSpeedHybridEngine的简单介绍可见下方；
3. 最终得到1个携有所有相关模型的对象rlhf_engine。

模型初始化的相关代码

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

  这里实际上是本框架最具特色的部分，除了定制有更详尽的DeepSpeed配置来进行更精细的优化管理外，DeepSpeed-Chat团队还专门为actor模型开发了一个名为“DeepSpeedHybridEngine”的优化引擎，正如其名“Hybrid（混合动力）”所述，由于actor在PPO训练的过程中，需要兼任训练（参数优化）与复杂推理（生成经验序列seq），普通的优化引擎只能胜任单一的训练优化或单一的推理优化，DeepSpeedHybridEngine将支持模型在两种模式中自动切换并享有相应优化，使得phase3的训练效率大幅提升，这也是DeepSpeed-Chat框架的优势所在。

3.3.3 根据prompt获取经验数据

3.3.3.1 经验数据获取过程

  再次借用ColossalChat的示意图来进行说明，经验数据的获取过程如下：

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（此reward非彼reward，图示的reward指InstructGPT所提及的“KL Reward”：为了防止对phase2学习到的reward过度自信，引入了SFT与logits的KL散度作为惩罚的Reward）等数据，在DeepSpeed-Chat中，于具体训练过程才开始计算。

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

# 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.logits2. reference:SFT策略-output_ref.logits3. reward:奖励-reward_score，InsructGPT中的r_\theta4. 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 = 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.3.2 seq的生成

  对于本次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 itcontinueelse: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.3.3 奖励reward_score和价值估计values的获取

  奖励模型的模型类RewardModel中实现了相应的方法forward_value()，可支持输入“一句对话”返回“环境奖励与价值估计”。与原先训练所用的方法forward()不同，forward()可支持输入“chosen-reject对话对”，主要实现了“对话对”之间排序损失的计算（forward()在【中篇】的2.3.3中已有所介绍，此处将不再赘述）。
  以下通过简单例子来对“奖励”以及“价值估计”作区分：

“奖励/环境奖励/reward_score”主要是为对话序列给出一个奖励值/做出评分，
“价值估计/values”是为对话序列中的每一个位置都给出价值预测，是与时间步/状态紧密相关的。有对话序列 seq=[11, 22, 33, 44, 55, 66, 0, 0, 0, 0]
其奖励reward_score只会是1个标量，如reward_score_seq=2.25；
其价值估计values是1维数组，如[2.01, 0.23, 2.89, 0.66, 0.33, 2.25, 0.36, 0.99, 1.32, 1.62]

# 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.2.3已经进行过详解，且与此处所述内容无关，此处不再赘述"""···def forward_value(···, return_value_only=False, ···):"""和forward有些差别，forward需要针对输入的chosen-rejected对计算排序损失并返回而forward_value只需要考虑一个输入，然后返回分值:param 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 valueselse:"""否则还将进一步取得reward_score"""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.3.4 策略模型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  = 3logits = [[[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.4 PPO训练数据管理-MiniDataset

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

1. Dataloader从Dataset中取出1个prompt_batch的无监督数据和1个prompt_batch的prompt数据；
2. 使用1个prompt_batch的prompt数据进行经验采集，将得到1个prompt_batch的经验数据；
3. 1个prompt_batch的无监督数据、1个prompt_batch的经验数据将被送入各自的MiniDataset实例进行管理：1个prompt_batch将被分成数个ppo_batch，供PPO训练进行数次迭代。

  上述第3步就是MiniDataset所要做的事，其函数方法分别执行了：

1. add()：获取batch（prompt_batch）数据；
2. seperate()：细分为ppo_batch数据；
3. free():清空获取到的batch数据并返回ppo_batch数据。

# 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)# 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_sizeself.small_batch_size = small_batch_sizedef seperate(self):"""维护1个small_dataset"""small_dataset = []# 从self.dataset中逐个取batchfor 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_datasetdef 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 Noneelse:raise ValueError("The dataset is full but we did not stop it. There is a bug in the code.")def free(self):"""清空self.dataset中的数据"""self.dataset = []

3.3.5 PPO训练过程

3.3.5.1 基本流程

  对于采集到的一批经验数据，使用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
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)

3.3.5.2 PPO训练

  1次PPO训练由train_rlhf()方法进行管理，其内部主要实现了：

1. KL散度惩罚奖励old_rewards的计算，为了防止对phase2学习到的环境奖励$r(x,y)$过度自信，所以加入了KL散度惩罚项：
   $$r_{KL}=r(x,y)-\beta \log \frac{{\pi }_{old}^{RL}(y|x)}{{\pi }^{SFT}(y|x)}$$
2. 优势advantages和回报returns的计算，其中包括本框架在内多数框架的advantages实现并非纯粹使用TD-error，而是在TD-error的基础上结合了MC方法，也即GAE（广义优势估计）；
   • 对于全长为$T$的轨迹来说，其某个时间步$t$的优势为（$\lambda =1$时，advantage完全使用MC方法；$\lambda =0$时，advantage完全使用TD-error方法）：
     $$\begin{gathered} {\hat{A}}_t={\delta }_t+(\gamma \lambda ){\delta }_{t+1}+(\gamma \lambda {)}^2{\delta }_{t+2}+\cdots +(\gamma \lambda {)}^{T-t+1}{\delta }_{T-1} \\ where{\delta }_t=r_{KL,t}+\gamma \cdot V_{old}(s_{t+1})-V_{old}(s_t) \end{gathered}$$
   • 回报returns就是奖励reward的累计，对于全长为$T$的轨迹来说，其到达某个时间步$t$时的回报为
     $$R_t={\hat{A}}_t+V_t$$
3. 在1个ppo_batch中，actor的损失计算公式为：
   $$pg\_loss=E_{\tau \sim {\pi }_{old}^{RL}}{E}_{(s_t,a_t)\sim \tau }[\max (-{\hat{A}}_t\cdot \frac{p_{new}^{RL}(a_t|s_t)}{p_{old}^{RL}(a_t|s_t)},-{\hat{A}}_t\cdot clip(\frac{p_{new}^{RL}(a_t|s_t)}{p_{old}^{RL}(a_t|s_t)},1-ϵ,1+ϵ))]$$
   其中，$\tau$指的仅是“answer”部分的内容，并不包括prompt部分。
4. 在1个ppo_batch中，critic的损失计算公式为：
   • 裁剪新价值估计$V_{new}$，使其不至于太偏离采集经验时的旧价值估计，使得经验回放仍能有效：
     $$V_{clip}=clip(V_{new},V_{old}-\varphi ,V_{old}+\varphi )$$
   • critic将拟合回报$R$：
     $$vf\_loss=\frac{1}{2}\cdot E_{\tau \sim {\pi }_{old}^{RL}}{E}_{s_t\sim \tau }[\max ((V_{new}(s_t)-R_t{)}^2,(V_{clip}(s_t)-R_t{)}^2)]$$
     其中，$\tau$指的仅是“answer”部分的内容，并不包括prompt部分。

  具体代码可见下方，为保证阅读的流畅性，我对其中的部分代码进行了调整，使得相应的函数代码衔接在其调用后方，便于具体对照其传参，从而辨析传入的新旧策略、新旧价值估计等：

# applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py
def train_rlhf(self, inputs):"""使用1个ppo_batch的经验数据，执行1次rlhf训练迭代"""# prompt input idsprompts = 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 idsseq = 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 = valueswith torch.no_grad():###计算KL惩罚修正的奖励################################################"""通过KL散度惩罚，以及r_\theta（来自reward model）计算得到修正的奖励，注意此处的入参：1. log_probs为经验数据中的旧策略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)"""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###计算优势与回报################################################"""计算优势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 = 0advantages_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"""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###计算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) * maskratio = torch.exp(log_ratio)"""计算策略梯度损失的2个情况：加权优势 与 裁剪加权优势"""pg_loss1 = -advantages * ratiopg_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()###计算critic损失并更新################################################"""将seq经验数据输入至critic，预测得到新价值估计"""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

3.3.5.3 无监督训练

  实际上就是常规的自回归语言建模任务。InstructGPT中提及，进行phase3的RLHF训练时，为使得模型在学习人类偏好的过程中仍能保有预训练模型解决任务的性能，引入了传统的自回归语言建模进行联合训练。
$$p(x)=\prod _{t=1}^{T}p(x_t|x_{<t})$$

# 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

3.3.5.4 EMA

  待完善…

3.4 实例测试

  “实例测试”与“指标评估”并不是完全相同的概念，实例测试是选择具体的数据实例输入进模型中，人工观察其输出结果，而非使用具体指标对结果进行评估。
  待完善…

3.5 相关拓展

3.5.1 phase3的参数设置

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

We have found that it is very unstable to use different generation training batch sizes (–per_device_train_batch_size) and PPO training batch sizes (–per_device_mini_batch_size), more than one PPO training epoch (–ppo_epochs), or more than one generation batch size (–generation_batch_numbers). These all point to the same problem: we are not able to update the actor model multiple times after generating experimental data. Therefore, in all of our successful runs, we have set per_device_train_batch_size = per_device_mini_batch_size and ppo_epochs = generation_batch_numbers = 1. This is unexpected for a standard RL training pipeline, and we have tried different methods to overcome this, but all have failed. One of the most likely reasons for this instability is that we found the log_probs and old_log_probs used in the actor_loss_fn function can quickly diverge even within two consecutive iterations, which causes the corresponding ratio to be huge. Setting a strict upper bound can alleviate this problem, but it cannot fully resolve the convergence issue.
We have also found that adding unsupervised training is not easy. We tried using the coefficient (–unsup_coef=27.8) provided by InstructGPT, but it caused instability in the RLHF training. According to InstructGPT, unsupervised training mainly affects the model quality on standard benchmarks instead of the RLHF performance. We did not put much effort into tuning this parameter.

3.5.2 PPO-ptx训练的迭代数对齐

  在phase3中，如果启用了无监督训练（PPO-ptx），那么无监督训练将是与PPO训练同步进行的，故两者的数据集处理几乎都是同步的，不仅是batch_size相同，以至于两者的batch数（step数）也都会被强制持平：通常情况下无监督数据量更大，按理同batch_size的情况下可迭代的次数也将多得多，但在PPO-ptx训练中，无监督训练的迭代数将会被裁至与PPO训练所用数据的迭代数持平——例如PPO所用训练数据满足迭代10次、无监督训练也只能进行10次迭代，多余的无监督数据将被弃用。

3.5.3 LMFlow的RAFT

  待完善…

3.6 版块相关问题

  暂无