instructGPT的前两阶段核心训练过程pytorch详细代码展示

训练细节

这篇内容主要是介绍关于instructGPT在训练的过程中代码细节。InstructGPT一共有三个训练阶段，分别是有监督的微调，reward模型的训练，以及PPO的训练。对于这三个阶段理论上有了之后，更加重要的是如何用代码来实现这些理论的细节。笔者认为，了解理论还不够，必须要真正的将理论用代码的方式实现出来，才是能真正的心安。

在以下的介绍中，会分别从数据的准备，模型的准备，和损失的计算三方面做各个阶段的代码介绍。注意，以下只是介绍核心的部分，从而了解核心后，读者可以自己应用到自己的框架中。

核心部分指的是对某一个小batch(1个或者多个样本)，演示如何准备必要的模型输入，模型的训练以及推理，以及模型的输出的损失计算。所以代码中不会包含Dataset类，train以及eval的函数，等通用的训练框架的构建。

1. 有监督的微调

在之前的介绍中，有监督的微调其实就是继续在initial_model上next token prediction的任务。

1.1 数据的准备

假设当前只有一个样本，这个样本是关于阿凡达电影的介绍，期望设计一个prompt从而，使用模型来生成对这个电影的评价，详细的batch的准备如下

raw_input = "《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。"
 
# 当前使用这个prompt，将其粘贴在raw的后面，从而读起来更加通顺
# 当前也可以设计多个prompt，和后面的label做笛卡尔积扩充数据也是没有问题的
prompt = "看完这部电影后，"
 
# label是在豆瓣上收集到的关于阿凡达的电影评价，当前之选3个样本，用于后期batch的构建
label = [
    "致喜欢但不怎么爱《阿凡达》的朋友：《阿凡达》让我大受震撼。这部续作在视觉效果、故事叙述和表演上都非常出色，让我在整个观影过程中都目瞪口呆",
    "感觉拥有我所见过的最令人难以置信的视觉特效，美得令人窒息（我看的3D版），但是故事本身比第一个要弱",
    "超出我的预期，大银幕属性极强，观赏性奇佳，亲情和对他人他物的关爱表达极易带入。"
]
 
# 通过for循环的方式将raw_input, prompt, 评价拼接起来
initial_data_ = [raw_input + prompt + i for i in label]

对于initial_model来说准备使用预训练过的GPT2的中文模型，uer/gpt2-chinese-cluecorpussmall，选择这个模型的原因主要是因为小，下载快且计算的速度快，便于结果的输出和验证。

from transformers import AutoTokenizer
 
initial_tokenizer = AutoTokenizer.from_pretrained('uer/gpt2-chinese-cluecorpussmall')

这里实例化的了一个gpt2用的tokenizer，使用这个tokenizer来对initial_data_做encode的动作。

# 先准备模型的输入
initial_token_info = initial_tokenizer.batch_encode_plus(
    initial_data_,
    max_length=256,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)
 
# 此时initial_token_info是一个字典，字典中包含'input_ids', 'token_type_ids', 'attention_mask'，这些keys正好是模型所需要的parameters。
# 更加重要的是当前使用了return_tensors='pt'， 自动的将所有的结果都转化成二纬的tensor数据，直接用于数据的计算

模型的输入已经准备完毕，但是对于labels还没有准备完毕。由于当前做的是next token prediction的任务，那么labels其实就是输入。注意，labels不用做shift的处理，将labels放入模型中之后，模型会自动对输入和输出做shift的动作，从而正确的计算损失。这里需要对labels做的是将padding的部分变成-100，这样最终模型自己算crossentropy的时候自动ignore掉-100上的损失。

以下以图的方式展示数据的处理方式，粉色的部分是raw_input, 蓝色的部分是prompt，黑色的部分是希望模型能生成的有监督的labels。

# 将labels也加到字典中，这是因为model的parameters中也有labels这个parameters
initial_token_info['labels'] = initial_token_info['input_ids']
initial_token_info['labels'] = torch.where(
    initial_token_info['labels']==0, 
    -100, 
    initial_token_info['labels']
)

这个时候关于initial阶段的输入部分也就准备完毕了，这里准备labels的方式是直接将padding的部分变成-100，从而来完成标准的next token prediction的任务的训练。当然其实也可以进一步严格一点，如果我们只关注模型在generated部分的损失，那么其实也可以将raw_input+prompt部分的整体的token也全部抹掉，具体做法如下。

注意以下代码全部注释，只是给出另外一种方案，不带入实际的计算

# initial_token_info['labels'] = initial_token_info['input_ids']
 
# 这里可以尝试重新encode下raw_input + prompt，并且不加任何的参数，得到list，算其length
# raw_prompt_len = len(initial_tokenizer.encode(raw_input + prompt))
 
# 因为每个样本的前面的部分都是raw_input+prompt, 都是相同的，所有将所有的行，前面相同的部分全部变成0
# initial_token_info['labels'][:, :raw_prompt_len] = 0
 
# 最终整体的来尝试将0的部分变成-100
# initial_token_info['labels'] = torch.where(
#     initial_token_info['labels']==0, 
#     -100, 
#     initial_token_info['labels']
# )

1.2 模型的训练

这个地方模型的训练还是非常的容易的，由于在之前的inital_token_info中已经准备好的所有模型所需要的输入了，并且initial_token_info字典的key和model的parameters名字一模一样，那么直接用**的方式即可传参数

# 首先看下initial_token_info包含哪些
# initial_token_info = {
#     'input_ids' : batch_size * max_length,
#   	'token_type_ids' : batch_size * max_length,
#   	'attention_mask' : batch_size * max_length,
#   	'labels' : batch_size * max_length		# 模型的输入中也可以传入labels, 会自动的去做shift以及teacher forcing, 并且使用cross_entropy_loss自动的算loss
# }
 
from transformers import AutoModelForCausalLM
initial_model = AutoModelForCausalLM.from_pretrained('uer/gpt2-chinese-cluecorpussmall')
outputs = initial_model(**initial_token_info)

1.3 损失的计算

由于模型中传了labels的值，模型自动算loss并且记录在ouputs中，取出即可

# 模型的ouputs自动带loss, loss拿出来做backward即可
outputs.loss.backward()

2 Reward函数的训练

在这个阶段中，主要的目标针对initial_model的prompt的输出，以及输出的打分数据，做reward模型的训练。这里的reward的选择有很多种。原paper中提到的方式是将initial_model最上层的embedding去掉，使用输入的last_hidden_state中的CLS token做句向量来做后续的计算，这种方式可行，但是主要是这种GPT的unidirectional模型计算有点慢，也可以直接用Bert模型来做也没有问题。在下面的代码中，reward模型选择使用bert-base-chinese来做。

在做之前，首先回顾一下之前reward的过程。首先是initial_model对某个prompt数据做generate的动作，比如生成了5个结果，然后再针对这5个结果，做打分的动作。

2.1 基于某一个prompt做结果生成

首先准备数据，这个数据则是raw_input + prompt

input_test_for_generate = raw_input + prompt
# 注意，这里一定要加return_tensors='pt'的动作，因为tokenizer会自动的将结果做成二维的tensor用于模型的推理
token_info = initial_tokenizer.encode_plus(
    input_test_for_generate,
    return_tensors='pt'
)
# 当前的token_info中已经记录下来最重要的input_ids, 在model.generate中只需要input_ids即可
input_ids = token_info['input_ids']
 
# 使用model.generate自动帮做beamsearch的动作
generated_result = initial_model.generate(
    input_ids=input_ids, 
    max_length=256, 
    do_sample=True, 		# 增加随机性，而非greedysearch
    num_beams=5, 
    num_return_sequences=5,   	# 每个样本生成5个结果
    no_repeat_ngram_size=3,	# 防止重复的token
    early_stopping=True		# 提前停止
)
 
# 使用batch_decode的方式对5个结果整体做decode的动作
generated_tokens = initial_tokenizer.batch_decode(
  generated_result, 
  skip_special_tokens=True	# decode之后自动去掉[CLS][SEP]那些token
)
 
# 生成之后token和token之间有空格，对5个结果的每个结果都去抹掉空格
generated_tokens = [i.replace(' ', '') for i in generated_tokens]

关于Transformer的生成模型的generate的方法，感兴趣的同学可以去查看https://huggingface.co/blog/how-to-generate

以下生成的结果

《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。看完这部电影后，我感觉到了一种奇妙的力量。这是一部非常好看的电影，也是我看过的最棒的一部动画电影。我觉得这不仅仅是电影本身的魅力，更是一种对于人性的深刻理解和对于生活的思考。一部优秀的动画片，能够让我们感受到人与人之间的关系，这种关系是多么的纯粹。我想，这是我最喜欢的动漫电影之一。
 
《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。看完这部电影后，我突然意识到，这是一部值得一看的电影，因为它给我们带来了新的思考，让我们看到一个真实的苏里。···影片讲述的是一个名叫苏里的女孩，在一次意外事故中，她遭遇了一场灾难，她的母亲不幸身亡，而她的父亲也因此受到了重大的伤害。
 
《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。看完这部电影后，我觉得这是一部值得一看的电影，因为它让我们看到了一个真实的苏里，一个关于爱情的真实故事。话说回来，电影中的人物很多，但是都有一个共同的特点，那就是他们都是一个人，他们有自己的生活，也有他们的家庭。
 
《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。看完这部电影后，我的脑海里浮现出一个画面：一个年轻的女孩子站在一个不起眼的地方，面对一个陌生的男人，她的内心是崩溃的，她不知道自己要做什么，但是她却一直在寻找答案。这个女孩是一个孤独的人，因为她没有人能够帮助她，所以她只能靠自己的力量来维持这个世界的秩序。在这个过程中，她遇到了一些困难，她想要找到解决问题的办法，这个男人帮助了她，让她找到了自己想要的东西。
 
《阿凡达》的背景设定在前一集的十年多以后，讲述苏里一家的故事、他们面临的麻烦、他们为了保护彼此而做的努力、他们为了存活而打的仗，以及他们所承受的悲剧。看完这部电影后，我觉得这个故事还是很有意思的，尤其是第一集，我们看到的是一个小男孩，他的爸爸是一名医生，妈妈是一位老师，还有一个孩子，他们的父亲是一对夫妻。他们在一起生活了一段时间，后来他们结婚了，结婚之后，孩子也跟着他们长大。在他们生活的这个过程中，父母的关系发生了很大的变化，这也是为什么他们之间的感情非常好的原因。因为他们是彼此的父母，而不是父母。

为了方便后期的训练，笔者通过手动的方式为这5个生成的结果做打分。这个时候，注意打分的方式有多种，这里暂时了解到的是两种打分方式，一种是assessment scores, 另外一种是preference scores。这两个名字是笔者自己取的。

2.2 打分方式

2.2.1 asessment scores

这种打分的代表数据是HFsummary，这个数据可以通过datasets的包去自动的load，有兴趣的读者可以自己去尝试。这种打分方式是针对单个样本做多维度的打分，比如对生成的每句话来说，这句话是否 [流畅，真实，和raw有关联]。以下是笔者对5句话的assessment打分，注意这里的打分不是特别的准确，只是举例子描述核心的训练过程而已。

这种数据适合的训练方式，可以是在Bert模型上方加个Linear层，从而用MAE或者MSE的方式来做损失的计算。

2.2.2 preference scores

这种打分的代表数据是WebGPT，这个数据也可以通过datasets的包去自动的load，有兴趣的读者可以自己去尝试。这种打分方式是单维度，是对生成的成对的样本来打分，判断在两个生成的样本中，到底哪个样本更加好一些。以下是笔者对5句话的preference打分，注意这里的打分是很原始的打分，后期会做成pairwise的形式从而来完成paper中所描述的损失的计算。

这种label可以使用两种方式的训练方式，这两种方式核心的思想差不多，但是损失函数的实现方式不一样

1. paper中pairwise的训练方式

2. 苏神所创建的cosent的训练方式

每一种训练的方式，都要对数据做不一样的准备，后面做详细的展示

2.3 Assessment打分下模型训练

2.3.1 数据集的准备

这种状态下的数据准备很简单的，因为不需要使用pairwise的方式，则简单的对数据做tokenizering的动作，后期放入模型即可

首先回顾下标签数据

scores_assessment = [
    [0.8, 0.2, 0.3],
    [0.7, 0.4, 0.7],
    [0.7, 0.3, 0.8],
    [0.75,0.4, 0.75],
    [0.6, 0.3, 0.3]
]

reward_tokenizer = AutoTokenizer.from_pretrained('bert-base-chinese')
 
# generate_tokens上面生成的token，这里统一改名字为reward_data
# 一定要注意的是，当前的reward_data里的5个样本，每个样本都是包含之前的raw_input和prompt的
# 这是因为希望模型不仅仅只关于generate的质量，也希望模型可以考虑generate的结果是否和raw_input相关
reward_data = generated_tokens
reward_token_info = reward_tokenizer.batch_encode_plus(
    reward_data,
    max_length=256,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)

2.3.2 模型的训练

模型可以是SequenceForClassification的模型，也就是在Bert的pooler的output上增加nn.Linear(hidden_size, 3)层。这种训练方式可以尝试使用图的方式做展示

从图中可以看到，BertForSequenceClassification模型是在Bert的上方加了一个Linear层的模型。当句子入模之后，在每个token上都会得到hidden向量，而[CLS]所对应的向量一般被认为是保存了句子中所有信息的向量。将这个向量拿出来，经过Bert的pooler层之后，再通过Linear层映射到3个维度上。此时的结果还需要经过Sigmoid，将值映射到0，1空间中，最终和标签参与损失的计算。以下是基于上图的代码实现。

# 使用SequenceClassification模型，这个模型是在base的模型的pool的结果上，多家了个linear层
from transformers import AutoModelForSequenceClassification
# nn.Linear(hidden_size, 3)的这个最后映射到多少个维度的值在num_labels上设置
reward_model = AutoModelForSequenceClassification.from_pretrained(
    'bert-base-chinese',
    num_labels=3
)
 
reward_outputs = reward_model(**reward_token_info)

2.3.3 损失的计算

import torch
to_sigmoid = torch.nn.Sigmoid()
 
# reward_outputs中的logits需要先做一个sigmoid映射到0，1的空间中，才能计算损失
reward_logits = to_sigmoid(reward_outputs.logits)
# 将labels从list转化成tensor
labels = torch.tensor(scores_assessment, dtype = torch.float)
# 使用MSELoss
loss_fn = torch.nn.MSELoss()    # 有的框架中使用L1Loss也就是mae的方式做的loss计算
# 或者使用MAEloss，以下注释
# loss_fn = torch.nn.L1Loss()
reward_loss = loss_fn(reward_logits, labels)
reward_loss.backward()

2.4 Preference打分下模型训练(paper的方式)

2.4.1 数据集的准备

首先需要尝试将数据做成pairwise的形式，从而的成对的来进行进行计算。回顾下标签的数据

scores_preference = [2,     1,     5,     3,     1]
#reward_data = 	    [sent1,  sent2, sent3, sent4, sent5]

对于这种标签其实一种最简单的方法还是像上面的介绍的，使用BertForSequenceClassification的方式来做，但是由于当前每个样本只有一个分数，num_labels需要设置成1，最终通过mse的损失来做模型的更新。

其实有更加不错的方法，也就是使用pairwise的方法，这种方法的思想来自simCSE。在将每个句子压缩一个维度(一个logits值)后，这个logits的值可以简单的理解成当前raw_input + prompt + generated整个句子的质量，这个值越大整个句子总体的质量越好。那么模型训练的过程中，可以尝试在5个句子中做两两句子的对比，在每个句子对中，期望质量高的句子的logits和质量低的句子的logits中间的差值被拉的越大越好。以下从图的角度来描述这一过程。

以下是详细的代码

# 将原数据做成以上的形态
batch_size = len(scores_preference)
ordered_reward_data = []
for i in range(batch_size):
    for j in range(batch_size):
        if scores_preference[i] > scores_preference[j]:
            ordered_reward_data += [generated_tokens[i], generated_tokens[j]]
 
# 假设当前的ordered_reward_data的个数有18个，那也就暗示着有9对数据
# 注意这个时候是没有label的，label的信息是隐藏在了ordered_reward_data的顺序中，也就是双数索引(从0开始)上的句子成对的要比单数索引上的分数高。
            
# 做分词的处理
reward_token_info = reward_tokenizer.batch_encode_plus(
    ordered_reward_data,
    max_length=256,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)

2.4.2 模型的训练

# 实例化模型, 注意这里num_labels为1，也就是我们只有一个维度的评价评分，将每个句子映射到一个维度上去
reward_model = AutoModelForSequenceClassification.from_pretrained(
    'bert-base-chinese',
    num_labels=1
)
 
# 生成logits
reward_outputs = reward_model(**reward_token_info)
# outputs中的logits的维度18 * 1。因为18个句子每个句子被最后的Linear映射到了1个维度上。

2.4.3 损失的计算

拿到outputs之后，由于准备数据的时候是，成对的数据是按顺序穿插起来的，为了方便计算损失，可以的展开。以下是损失计算的方式

# 整理成句子对
reward_logits = reward_outputs.logits.reshape(-1, 2)
# 这个时候reward_logits的维度为18 * 1, reshape后为9 * 2
# 通过reshape的动作之后，将双单的向量矩阵分开，放在了第一列和第二列，成对的来看，第一列的句子的分数要比第二列的高
 
# 设计论文中所提到的RankLoss
class RankLoss(torch.nn.Module):
    def __init__(self, eps=1e-8) -> None:
        super().__init__()
        self.eps = eps
        self.log_sigmoid = torch.nn.LogSigmoid()
 
    def forward(self, pos, neg):
      	# 注意这里传进来的pos，neg的向量的维度都是batch, 1
        # pos,neg是成对的向量，在每一个pair中，前面的句子比后面的句子分数要高
        # 下面的log_sigmoid的目标就是尝试拉长这个每一对中前者何后者之前的差距
        # 从而让模型学习到前者的结果要优于后者
        # 这个时候为什么没有labels呢？这是因为labels的信息已经被隐藏在pos,neg之前穿插的数据顺序中
        return -self.log_sigmoid(pos - neg + self.eps).mean()
      
loss_fn = RankLoss()
reward_loss = loss_fn(reward_logits[:, 0], reward_logits[:, 1])
reward_loss.backward()

2.5 Preference打分下模型训练(consent的方式)

在2.4中，是将raw_input + prompt + generated 三部分全部添加在一起做成了一个整体的句子，再通过句子对的方式来比较。是否可以有更加灵活一点的方法呢，比如将raw,prompt和generated也做成句子对，从而通过拉大“句子对”的对之间的差距从而更新模型呢？

笔者的这种不同于paper里的思路，是从苏神的cosent训练词向量的方法来的，关于cosent的损失的细节可参考https://spaces.ac.cn/archives/8847。笔者认为，生成generated质量越高，则其和raw + prompt的相似度就应该越高(或者使用新的Linear层u, v处理下也可以)，从而将其做成句子对来对"句子对"对训练更加合理。另外一点好处的是在苏神的文章中提到，cosent的优势就是更加适合有排序性的标签，也就是当前的这种Preference的分数。以下给出一个大致的处理图

2.5.1 数据的准备

期望做成以下的结果

raw, prompt + sent1
raw, prompt + sent2
raw, prompt + sent3
raw, prompt + sent4
raw, prompt + sent5

这里可以减少计算，先做成

raw
prompt + sent1
prompt + sent2
prompt + sent3
prompt + sent4
prompt + sent5

得到各个句子的CLS向量后，再将raw_prompt给提取出来，复制5份做成以上期望的样本

# 将原generated_tokens里面的raw去掉
generated_tokens_only_generated = [i[len(raw_input):] for i in generated_tokens]
# 然后list的元素的第一个位置增加raw
generated_tokens_only_generated = [raw_input] + generated_tokens_only_generated
 
# 做分词的动作
reward_token_info = reward_tokenizer.batch_encode_plus(
    generated_tokens_only_generated,
    max_length=256,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)

2.5.2 模型的训练

# 其实可以直接复用上面的reward_model, 这里为了完整性，重新实例化一遍
reward_model = AutoModelForSequenceClassification.from_pretrained(
    'bert-base-chinese',
    num_labels=1
)
 
# 我们不需要reward_model的linear的输出了，想要其base模型
raw_model_base = reward_model.base_model
 
# 在ouputs中，我们关注的是last_hidden_state的CLS的部分
outputs = raw_model_base(**reward_token_info)

2.5.3 损失的计算

# 首先拿出last_hidden出来，维度为6 * seq_length * 768, 取CLS的部分
cls_hiddens = outputs.last_hidden_state[:, 0, :]
 
# 6个是因为第一个是raw_input的CLS，而后面的5个是prompt+5个generate的CLS
raw_input_hidden_states = cls_hiddens[0]
 
# 取出raw_input的CLS，维度为768, 将复制5份，做成5 * 768的维度
raw_input_hidden_states = raw_input_hidden_states.repeat(5, 1)
 
# 起初generated的部分，维度也为5 * 768, 因为有5个句子
generated_token_hidden_states = cls_hiddens[1:]
 
# 在cosent计算的时候labels的信息需要有，是用来看到底哪个句子对分数高，哪个句子对分数低用的
labels = torch.tensor(scores_preference)
 
# 这样维度一样，可以做后面句子对的cosent的损失计算
# 设计关于cosent的损失函数
def cosent_task_function(cls1, cls2, labels):
    cossim_for_each_pair = (cls1 * cls2).sum(axis=1)
    cossim_ij_diff_matrix = cossim_for_each_pair[:, None] - cossim_for_each_pair[None, :]
    labels_ij_diff_matrix = labels[:, None] - labels[None, :]
    labels_ij_diff_matrix = (labels_ij_diff_matrix < 0).to(cossim_for_each_pair.device, dtype=torch.float)
    cossim_ij_diff_matrix = cossim_ij_diff_matrix - (1 - labels_ij_diff_matrix) * 1e12
    cossim_ij_diff_matrix = cossim_ij_diff_matrix.reshape(-1)
    cossim_ij_diff_matrix = torch.cat(
        [
            torch.tensor([0]).to(cossim_ij_diff_matrix.device, dtype=torch.float),
            cossim_ij_diff_matrix
        ], dim=0
    )        
    loss = torch.logsumexp(cossim_ij_diff_matrix, dim=0)
    loss = loss / cls2.shape[0]
    return loss
 
loss = cosent_task_function(
  raw_input_hidden_states, 
  generated_token_hidden_states, 
  labels
)
loss.backward()

2.6 Assessment打分下的Cosent

Cosent的方法如何用刀Assessment的打分呢？其实很简单，Preference打分可以看成是用户在一个维度上的打分，Assessment有n个维度，那么可以尝试设计n个维度上的Cosent损失函数，TotalLoss = a1 * CosentLoss1(真实度) + a2 * CosentLoss2(准确度) + a3 * CosentLoss3(相关度) + a4 * (无敏感词度)从而完成Loss的计算，注意a1 + a2 + a3 + a4 = 1, 如果想让模型更加的关注无敏感词这个维度，可以手动的增加a4的值，从而让模型更加关注这个维度上的信息。