大模型入门（六）—— RLHF微调大模型

一、RLHF微调三阶段

参考：https://huggingface.co/blog/rlhf

1）使用监督数据微调语言模型，和fine-tuning一致。

2）训练奖励模型

奖励模型是输入一个文本序列，模型给出符合人类偏好的奖励数值，这个奖励数值对于后面的强化学习训练非常重要。构建奖励模型的训练数据一般是同一个数据用不同的语言模型生成结果，然后人工打分。如果是训练自己领域的RLHF模型，也可以尝试用chatgpt打分，效果也不错。

3）训练RL模型

在训练强化学习模型时，需要搞清楚状态空间、动作空间、策略函数、价值函数这些东西，动作空间就是所有的token，状态空间就是输入的序列的分布，价值函数由第二步的奖励模型和策略约束结合，策略函数就是微调的大模型。

从上图可以看出，给定一个输入xx，会生成两个文本y1y1和y2y2，一个来自于初始的模型，另一个来自于微调的模型，微调的模型生成的文本还会进入到奖励模型中打分输出rθrθ，而初始模型和微调的模型生成的结果会用KL散度约束它们的分布，确保模型不会太偏离原来的模型，并且能输出高质量的回复。

值得注意的是三个阶段的训练数据尽量是分布一致的，否则后面的训练会很不稳定。所以在第一步微调时不要一味地使用大量的训练数据（这一步的数据比较容易获得），尽量和后面两步的数据分布保持一致。

二：RLHF代码理解

以DeepSpeed-Chat的代码去理解RLHF的过程。https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat

DeepSpeed-Chat提供了RLHF三个阶段的训练代码，可以很方便地训练三个阶段，现在我们来一个一个阶段地来看。

1）数据集处理

首先从数据集的处理出发，去理解三个阶段的输入是什么样的数据？在training/utils/data/raw_datasets.py提供了多种开源数据集的读取方式，可以看到每个数据集都包含prompt（提问），chosen（正向回答），rejected（负向回答）。以其中某一个为例：

class DahoasRmstaticDataset(PromptRawDataset):

    def __init__(self, output_path, seed, local_rank, dataset_name):
        super().__init__(output_path, seed, local_rank, dataset_name)
        self.dataset_name = "Dahoas/rm-static"
        self.dataset_name_clean = "Dahoas_rm_static"

    def get_train_data(self):
        return self.raw_datasets["train"]

    def get_eval_data(self):
        return self.raw_datasets["test"]

    def get_prompt(self, sample):
        return sample['prompt']

    def get_chosen(self, sample):
        return sample['chosen']

    def get_rejected(self, sample):
        return sample['rejected']

    def get_prompt_and_chosen(self, sample):
        return sample['prompt'] + sample['chosen']

    def get_prompt_and_rejected(self, sample):
        return sample['prompt'] + sample['rejected']

具体的数据处理在training/utils/data/data_utils.py中，下面的代码展示了三个阶段使用的输入是什么？在第一步，即监督微调大模型，使用prompt + chosen；在第二步，即训练奖励模型时，需要使用prompt + chosen 和 prompt + rejected；在第三步，即训练RL模型，只使用prompt。

if train_phase == 1:
        for i, tmp_data in enumerate(current_dataset):
            # tokenize the text
            chosen_sentence = raw_dataset.get_prompt_and_chosen(
                tmp_data)  # the accept response
            if chosen_sentence is not None:
                chosen_sentence += end_of_conversation_token
                chosen_token = tokenizer(chosen_sentence,
                                         max_length=max_seq_len,
                                         padding="max_length",
                                         truncation=True,
                                         return_tensors="pt")
                chosen_token["input_ids"] = chosen_token["input_ids"].squeeze(
                    0)
                chosen_token["attention_mask"] = chosen_token[
                    "attention_mask"].squeeze(0)
                chosen_dataset.append(chosen_token)

    elif train_phase == 2:
        for i, tmp_data in enumerate(current_dataset):
            # tokenize the text
            c