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

一、RLHF微调三阶段

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

2）训练奖励模型

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

3）训练RL模型

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

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

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

二：RLHF代码理解

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'\]

下面的代码展示了三个阶段使用的输入是什么？在第一步，即监督微调大模型，使用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
            chosen\_sentence = raw\_dataset.get\_prompt\_and\_chosen(
                tmp\_data)  # the accept response
            reject\_sentence = raw\_dataset.get\_prompt\_and\_rejected(
                tmp\_data)  # the accept response
            if chosen\_sentence is not None and reject\_sentence is not None:
                chosen\_sentence += end\_of\_conversation\_token  # the accept response
                reject\_sentence += end\_of\_conversation\_token
                chosen\_token \= tokenizer(chosen\_sentence,
                                         max\_length\=max\_seq\_len,
                                         padding\="max\_length",
                                         truncation\=True,
                                         return\_tensors\="pt")
                reject\_token \= tokenizer(reject\_sentence,
                                         max\_length\=max\_seq\_len,
                                         padding\="max\_length",
                                         truncation\=True,
                                         return\_tensors\="pt")
                chosen\_token\["input\_ids"\] = chosen\_token\["input\_ids"\]
                chosen\_token\["attention\_mask"\] = chosen\_token\["attention\_mask"\]
                chosen\_dataset.append(chosen\_token)

                reject\_token\["input\_ids"\] = reject\_token\["input\_ids"\]
                reject\_token\["attention\_mask"\] = reject\_token\["attention\_mask"\]
                reject\_dataset.append(reject\_token)

    elif train\_phase == 3:
        for i, tmp\_data in enumerate(current\_dataset):
            # tokenize the text
            prompt = raw\_dataset.get\_prompt(tmp\_data)
            if prompt is not None:
                prompt\_token \= tokenizer(prompt, return\_tensors="pt")
                prompt\_token\["input\_ids"\] = prompt\_token\["input\_ids"\]
                prompt\_token\["attention\_mask"\] = prompt\_token\["attention\_mask"\]
                for key\_word in \["input\_ids", "attention\_mask"\]:
                    length \= prompt\_token\[key\_word\].size()\[-1\]
                    if length > max\_seq\_len:
                        y \= prompt\_token\[key\_word\].squeeze(0)\[length -
                                                              (max\_seq\_len \-
                                                               1):\].flip(0)
                    else:
                        y \= prompt\_token\[key\_word\].squeeze(0).flip(0)
                    prompt\_token\[key\_word\] \= y
                prompt\_dataset.append(prompt\_token)

2）监督数据微调大模型

文件夹下的traning_scripts提供了单GPU，多GPU，多机器的训练脚本。监督数据微调没有什么值得说的，和我们常用的微调方式一致，可以选择一个开源的模型，比如facebook/opt-1.3b，微调时可以选择lora和offload微调。

3）训练奖励模型

奖励模型可以选择一个较小的模型，如opt-350M，在chosen和rejected这种样本对上训练。reward model的输出类似于回归任务，将大模型的输出，然后经过N ✖️ 1 的线性层，得到一个batch size ✖️ seq len ✖️ 1的输出。在训练过程中，使用到的loss是二元交叉熵，确保每个prompt 的 chosen分数都是要大于rejected。

loss += -torch.log(torch.sigmoid(c\_truncated\_reward \- r\_truncated\_reward)).mean()

上面的代码中c_truncated_reward 和 r_truncated_reward 即给定一个prompt，对应的chosen和rejected获得的分数，而且是chosen 和 rejected所有token的分数差值。注意在这里因为chosen和rejected的长度不一致，而且还有padding的部分，所以c_truncated_reward和r_truncated_reward要做阶段，主要是截取chosen_id和rejected_id不等的部分出来，去除共同padding的部分。

4）训练强化学习模型

在第三步我们需要两个模型，一个是第一步训练好的SFT模型，另一个是第二步训练好的reward模型。接下来看下强化学习模型训练的步骤：

1）初始化rlhf engine，在代码training/step3_rlhf_finetuning/main.py中

rlhf\_engine = DeepSpeedRLHFEngine(
        actor\_model\_name\_or\_path\=args.actor\_model\_name\_or\_path,
        critic\_model\_name\_or\_path\=args.critic\_model\_name\_or\_path,
        tokenizer\=tokenizer,
        num\_total\_iters\=num\_total\_iters,
        args\=args)

rlhf_engine中会包含4个模型对象：self.actor: sft模型，可训练，作为策略模型；self.ref: sft模型，不可训练，只做前向推断，用于约束self.actor生成结果的向量分布；self.critic：reward模型，可训练，价值模型，用于对生成的每个动作打分；self.reward：reward 模型，不可训练，用于计算整个序列的奖励值。

2）初始化ppo训练器，在代码training/step3_rlhf_finetuning/main.py中

ppo\_trainer = DeepSpeedPPOTrainerUnsupervised if unsupervised\_training\_enabled else DeepSpeedPPOTrainer
trainer \= ppo\_trainer(rlhf\_engine, args)

3）生成PPO的训练样本，在代码training/step3_rlhf_finetuning/main.py中

out = trainer.generate\_experience(prompts)
exp\_dataset \= exp\_mini\_dataset.add(out)

4）训练PPO模型，在代码training/step3_rlhf_finetuning/main.py中

for ppo\_ep in range(args.ppo\_epochs):
    for i, (exp\_data, unsup\_data) in enumerate(
            zip(exp\_dataset, unsup\_dataset)):
        actor\_loss, critic\_loss \= trainer.train\_rlhf(exp\_data)                                

回过头来再看rlhf_engine和ppo_trainer的实现逻辑。

rlhf_engine的实现在training/step3_rlhf_finetuning/rlhf_engine.py中，主要是初始化了几个模型对象

self.actor = self.\_init\_actor(
    actor\_model\_name\_or\_path\=actor\_model\_name\_or\_path)
self.ref \= self.\_init\_ref(
    actor\_model\_name\_or\_path\=actor\_model\_name\_or\_path)
self.actor\_ema \= None
if self.args.enable\_ema:
    self.actor\_ema \= self.\_init\_ema(
        actor\_model\_name\_or\_path\=actor\_model\_name\_or\_path)

self.critic \= self.\_init\_critic(
    critic\_model\_name\_or\_path\=critic\_model\_name\_or\_path)
self.reward \= self.\_init\_reward(
    critic\_model\_name\_or\_path\=critic\_model\_name\_or\_path)

ppo_trainer的具体代码在training/step3_rlhf_finetuning/ppo_trainer.py中，由于对强化学习不是很熟悉，只能简单地描述下整个逻辑：

1）输入prompt，使用self.actor生成对应的answer，并拼接成一个完整的seq，这其实是一个采样的过程，类似于强化学习中生成一条完整的状态-动作序列。动作即生成的token，状态时生成token的前缀输入。

with torch.no\_grad():
    seq \= self.actor\_model.module.generate(prompts,
                                            max\_length\=max\_min\_length,
                                            min\_length\=max\_min\_length)

2）基于当前 T 时刻的网络参数，生成完整的状态-动作序列，奖励值。

with torch.no\_grad():
    output \= self.actor\_model(seq, attention\_mask=attention\_mask)
    output\_ref \= self.ref\_model(seq, attention\_mask=attention\_mask)
    reward\_score \= self.reward\_model.forward\_value(
        seq, attention\_mask,
        prompt\_length\=self.prompt\_length)\['chosen\_end\_scores'\].detach(
        )
    values \= self.critic\_model.forward\_value(
        seq, attention\_mask, return\_value\_only\=True).detach()\[:, :-1\]

logits \= output.logits
logits\_ref \= output\_ref.logits

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）训练PPO模型，所有的核心训练逻辑都在这里。

def train\_rlhf(self, inputs):
    # train the rlhf mode here
    #\## process the old outputs
    prompts = inputs\['prompts'\]
    log\_probs \= inputs\['logprobs'\]
    ref\_log\_probs \= inputs\['ref\_logprobs'\]
    reward\_score \= inputs\['rewards'\]
    values \= inputs\['value'\]
    attention\_mask \= inputs\['attention\_mask'\]
    seq \= inputs\['input\_ids'\]

    start \= prompts.size()\[-1\] - 1
    action\_mask \= attention\_mask\[:, 1:\]

    old\_values \= values
    with torch.no\_grad():  
　　　　# 1、计算生成的每个token的奖励值
        old\_rewards \= self.compute\_rewards(prompts, log\_probs,
                                            ref\_log\_probs, reward\_score,
                                            action\_mask)  
　　　　# 2、计算价值，价值不等于奖励值，价值是考虑到未来的，奖励值只考虑当下
        advantages, returns \= self.get\_advantages\_and\_returns(
            old\_values, old\_rewards, start)

    #\## process the new outputs
    batch = {'input\_ids': seq, "attention\_mask": attention\_mask}
    actor\_prob \= self.actor\_model(\*\*batch, use\_cache=False).logits
    actor\_log\_prob \= gather\_log\_probs(actor\_prob\[:, :-1, :\], seq\[:, 1:\])  
　　# 3、计算actor网络的loss，并更新网络参数
    actor\_loss \= self.actor\_loss\_fn(actor\_log\_prob\[:, start:\],
                                    log\_probs\[:, start:\], advantages,
                                    action\_mask\[:, start:\])
    self.actor\_model.backward(actor\_loss)
    self.actor\_model.step()
    value \= self.critic\_model.forward\_value(\*\*batch,
                                            return\_value\_only\=True,
                                            use\_cache\=False)\[:, :-1\]  
　　# 4、计算critic网络的loss，并更新网络参数
    critic\_loss \= self.critic\_loss\_fn(value\[:, start:\], old\_values\[:,
                                                                    start:\],
                                        returns, action\_mask\[:, start:\])
    self.critic\_model.backward(critic\_loss)
    self.critic\_model.step()

    return actor\_loss, critic\_loss

因为这一部分的内容比较多，我们再细分来描述：

3.1）计算每个时刻（沿着序列的方向定义时刻）的奖励值，即给定前缀输入，生成当前token时对应的奖励值，奖励值由两部分组成，一是完整的序列奖励，由self.reward输出的，二是self.actor和self.ref输出的token向量的KL散度值。具体的代码：

def compute\_rewards(self, prompts, log\_probs, ref\_log\_probs, reward\_score,
                    action\_mask):

    kl\_divergence\_estimate \= -self.kl\_ctl \* (log\_probs - ref\_log\_probs)
    rewards \= kl\_divergence\_estimate
    start \= prompts.shape\[1\] - 1
    ends \= start + action\_mask\[:, start:\].sum(1)
    reward\_clip \= torch.clamp(reward\_score, -self.clip\_reward\_value,
                                self.clip\_reward\_value)
    batch\_size \= log\_probs.shape\[0\]
    for j in range(batch\_size):
        rewards\[j, start:ends\[j\]\]\[\-1\] += reward\_clip\[j\]

    return rewards

3.2）计算每个时刻的价值，actor的价值采用TD误差。在这里要指明价值不等于奖励值，奖励值只取决于当前时刻的状态和动作，而价值是考虑到了未来的情况的。所以价值的计算如下：核心是下面的delta的计算，除了考虑到当前的时刻的奖励值，还考虑到了未来时刻的输出的奖励值nextvalues（只不过这里的奖励值是由critic网络直接输出的每个token对应的分数）。下面的函数输出了两个值，一个是advantages，用于更新actor。二是returns，这是我们的目标Q值，用于后面更新critic。

def get\_advantages\_and\_returns(self, values, rewards, start):
    # Adopted from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling\_ppo.py#L134
    lastgaelam = 0
    advantages\_reversed \= \[\]
    length \= rewards.size()\[-1\]
    for t in reversed(range(start, length)):
        nextvalues \= values\[:, t + 1\] if t < length - 1 else 0.0
        delta \= rewards\[:, t\] + self.gamma \* nextvalues - values\[:, t\]  # TD误差
        lastgaelam = delta + self.gamma \* self.lam \* lastgaelam
        advantages\_reversed.append(lastgaelam)
    advantages \= torch.stack(advantages\_reversed\[::-1\], dim=1)
    returns \= advantages + values\[:, start:\]  # Q值
    return advantages.detach(), returns

3.3）计算actor的loss，其实在这里有一个目标actor和当前actor的概念，下面函数中old_logprobs是由目标actor输出的，logprobs是由当前actor输出的，ppo算法中actor的loss是使用当前actor输出的概率和目标actor输出的概率的比值来约束的，转换成log后就是logprobs - old_logprobs。

def actor\_loss\_fn(self, logprobs, old\_logprobs, advantages, mask):
    #\# policy gradient loss
    log\_ratio = (logprobs - old\_logprobs) \* mask  # 当前策略和上一策略的比值
    ratio = torch.exp(log\_ratio)
    pg\_loss1 \= -advantages \* ratio
    pg\_loss2 \= -advantages \* torch.clamp(ratio, 1.0 - self.cliprange,
                                            1.0 + self.cliprange)
    pg\_loss \= torch.sum(torch.max(pg\_loss1, pg\_loss2) \* mask) / mask.sum()
    return pg\_loss

3.4）计算critic的loss，在这里一样有一个目标critic和当前cirtic的概念，用当前cirtic输出的value减去目标Q值（也就是上面计算得到的returns）再求平方，所以是使用了当前critic的Q值和目标critic的Q值的均方误差作为critic的loss。

def critic\_loss\_fn(self, values, old\_values, returns, mask):
    #\# value loss
    values\_clipped = torch.clamp(
        values,
        old\_values \- self.cliprange\_value,
        old\_values + self.cliprange\_value,
    )
    vf\_loss1 \= (values - returns)\*\*2  # 当前critic和目标critic的Q值的均方误差
    vf\_loss2 \= (values\_clipped - returns)\*\*2
    vf\_loss \= 0.5 \* torch.sum(
        torch.max(vf\_loss1, vf\_loss2) \* mask) / mask.sum()
    return vf\_loss

在上面3.3）和3.4）提到了目标actor和目标critic，但在代码里并没有创建这两个变量，这里其实和PPO的训练方式有关，首先是利用T时刻的actor和critic生成状态-动作序列和价值，奖励等并存储下来。在训练PPO时会使用生成的状态-动作去重新输出价值、奖励等值并更新actor和critic参数，所以并没有显示构造目标actor和目标critic，但是存储了它们产生的结果。存储的这部分数据会不断地更新，保证目标actor和critic和当前的actor和critic的参数不会有太大的差别，更新的逻辑在training/utils/data/data_utils.py中的MiniDataset类中。

所以强化学习模型的训练流程就是两步，一是先生成目标actor和critic的值作为对比的数据，二是训练actor和critic模型，将代码简化，其实就是training/step3_rlhf_finetuning/main.py中下面的代码段：

for epoch in range(args.num\_train\_epochs):
    ....
    for step, (batch\_prompt, batch\_unsupervised) in enumerate(
            zip(prompt\_train\_dataloader, unsupervised\_train\_dataloader)):
        ...

        # 生成训练强化学习模型的数据
        out = trainer.generate\_experience(prompts)
        exp\_dataset \= exp\_mini\_dataset.add(out)

        if exp\_dataset is not None:
            ...
            # 训练强化学习模型
            for ppo\_ep in range(args.ppo\_epochs):
                for i, (exp\_data, unsup\_data) in enumerate(
                        zip(exp\_dataset, unsup\_dataset)):
                    actor\_loss, critic\_loss \= trainer.train\_rlhf(exp\_data)
                    actor\_loss\_sum += actor\_loss.item()
                    critic\_loss\_sum += critic\_loss.item()
                    average\_reward += exp\_data\["rewards"\].mean()

最后的最后

感谢你们的阅读和喜欢，我收藏了很多技术干货，可以共享给喜欢我文章的朋友们，如果你肯花时间沉下心去学习，它们一定能帮到你。

因为这个行业不同于其他行业，知识体系实在是过于庞大，知识更新也非常快。作为一个普通人，无法全部学完，所以我们在提升技术的时候，首先需要明确一个目标，然后制定好完整的计划，同时找到好的学习方法，这样才能更快的提升自己。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。

四、AI大模型商业化落地方案

五、面试资料

我们学习AI大模型必然是想找到高薪的工作，下面这些面试题都是总结当前最新、最热、最高频的面试题，并且每道题都有详细的答案，面试前刷完这套面试题资料，小小offer，不在话下。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】