从零实现强化学习RLHF代码（PPO、RLOO）

从零实现RLHF

强化学习的过程也叫做偏好对齐的过程，发展到现在其已不再神秘了，一些最新的大模型流程基本都有RLHF的过程，之前看过一篇论文其是先DPO然后再PPO，DPO的文章之前我们已经讲过了，所以现在我们以PPO为例，进行RLHF全过程代码的构建。

所有代码已集成至git仓库：https://github.com/mst272/LLM-Dojo/tree/main/rlhf，这是一个强化学习轻量级框架，使用简洁的代码对各种强化学习方法进行了集成，支持deepspeed单卡或多卡训练，且便于自己修改与学习。除了PPO之外，还有RLOO、SimPO等等方法。

本篇文章着重讲解具体的代码实现过程，其中的原理一篇文章是讲不完的，需要读者自己去阅读及思考。

首先，关键点就在于三个步骤

SFT就是你微调过后，想要RLHF的模型，也就是想要提升的模型。
剩余两个在下面进行介绍

Reward

Reward model

Reward model实际上是一个分类器，对输入的 指令+回答 进行 打分。一般Reward model我们就采用常规的已经pretrain的分类模型即可，也可以自己训练魔改一个。例如可以使用gpt2、facebook/opt-350m等，加载方式如下即可，如果 AutoModelForSequenceClassification不能识别的话需要自己在模型的config.json文件中进行映射了：
model = AutoModelForSequenceClassification.from_pretrained("gpt2")

如果想使用大模型的话，例如QWEN、LLAMA等模型，可以直接通过huggingface使用ForSequenceClassification方法加载，可以去transformers库中查看，一般的大模型都会有一个ForSequenceClassification类，我们以最新的QWEN2为例，进入transformers库中的modeling_qwen2.py， 如下图所示，可以看到其是有一个Qwen2ForSequenceClassification方法的，其他模型也以此类推，使用这个方法作为reward model的加载。

再进一步，我们看到self.score是一个线性层被添加到了QWEN模型中。进一步看下面代码：

经过一个线性层后的logits维度为：[batch_size, sequence_length, num_classes]，经过图上第二个红框的操作后logits维度变为：[batch_size, num_classes]，这里实现的逻辑是对于每句话，选最后一个词的得分作为这句话的得分，以供后续reward 训练。

由于模型大小原因，我们本次以小一点的Phi-3为例进行reward的训练
代码示例：

from transformers import Phi3ForSequenceClassification  # 不同模型需要替换不同接口
model = Phi3ForSequenceClassification.from_pretrained(
    model_config.model_name_or_path, num_labels=1, **model_kwargs
)

另外一种方法

如果你觉得大模型作为reward model太大的话，也可以通过config进行设置，以减少模型的层数等，看自己偏好了，以下是一个代码示例：

import torch
import torch.nn.functional as F
from transformers import LlamaModel, LlamaConfig, LlamaForCausalLM, LlamaForSequenceClassification

torch.manual_seed(1)

# 本地创建模型
config = LlamaConfig(vocab_size = 100,      # default is 32000
                    hidden_size = 256,
                    intermediate_size = 512,
                    num_hidden_layers = 2,
                    num_attention_heads = 4,
                    num_key_value_heads = 4,
                    )
model = LlamaForCausalLM(config)
model.save_pretrained('./lm_pretrained')

# load LLaMA with "one labels" sequence classification model
rm_model = LlamaForSequenceClassification.from_pretrained('./lm_pretrained', num_labels=1)

print(rm_model)

通过上述代码减少模型层数后，又可以作为一个新的reward model了。

Reward Train

使用hf的trl库我们可以很容易的进行Reward Train的训练工作。我们以huggingface的Anthropic/hh-rlhf数据集为例：

使用其他数据集的话将其改成这种格式即可。这里我们并没有适配每个模型的chat template，如果想做的话可以参考上述git项目中的做法，为了简化流程我这里就直接用这个数据集进行下一步操作了。

RewardTrainer 接受的数据需要有非常特定的格式，需要包括：

所以我们需要先进行数据处理操作，代码放在下面，也是比较简单了，因为现在的库真的是越来越好用了：

    def preprocess_function(examples):
        new_examples = {
            "input_ids_chosen": [],
            "attention_mask_chosen": [],
            "input_ids_rejected": [],
            "attention_mask_rejected": [],
        }
        for chosen, rejected in zip(examples["chosen"], examples["rejected"]):
            tokenized_chosen = tokenizer(chosen)
            tokenized_rejected = tokenizer(rejected)

            new_examples["input_ids_chosen"].append(tokenized_chosen["input_ids"])
            new_examples["attention_mask_chosen"].append(tokenized_chosen["attention_mask"])
            new_examples["input_ids_rejected"].append(tokenized_rejected["input_ids"])
            new_examples["attention_mask_rejected"].append(tokenized_rejected["attention_mask"])

        return new_examples

有需要的话可以进行一下长度截断等过滤，以及划分一下train 和 test。
最后在配置一下相关的参数即可开始训练了。

然后将数据集送入Trainer开始训练.

trainer = RewardTrainer(
        model=model,
        tokenizer=tokenizer,
        args=config,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=get_peft_config(model_config),
    )
    trainer.train()
    trainer.save_model(config.output_dir)
    metrics = trainer.evaluate()
    trainer.log_metrics("eval", metrics)
    print(metrics)

RL (Reinforcement Learning): PPO 、RLOO

完成奖励模型的训练后我们来到最后一步，强化学习阶段。需要用到上一步的 Reward模型和原始的SFT模型。

训练用到的PPO Trainer依然是来自我们的huggingface trl库，可以很方便的使用及训练。目前最新版本是PPOv2Trainer，其是之前的PPOTrainer的改进版本，作者原话如下：

The PPOv2Trainer is the new experimental PPO trainer we now recommend to the users. It’s a refactor of PPOTrainer and PPOv2Trainer introduces more uniform APIs, better logging, documentations, and more benchmark results.

这里用到的数据集只需要prompt部分，所以应该跟上述的Reward Model训练的数据集相吻合（只用到了chosen和reject部分）。

首先需要加载四个模型：分别是value_model、reward_model、ref_policy、policy。

value_model = AutoModelForSequenceClassification.from_pretrained(
        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
    )
    reward_model = AutoModelForSequenceClassification.from_pretrained(
        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
    )
    ref_policy = AutoModelForCausalLM.from_pretrained(
        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
    )
    policy = AutoModelForCausalLM.from_pretrained(
        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
    )

其中我们要训练的就是policy模型，加载完模型后还需要进行数据处理操作，与上述reward训练时有所不同，RL训练只需要返回prompt编码后的input_ids即可。

 def prepare_dataset(dataset, tokenizer):
        """pre-tokenize the dataset before training; only collate during training"""

        def tokenize(element):
            outputs = tokenizer(
                element[dataset_text_field],
                padding=False,
            )
            return {"input_ids": outputs["input_ids"]}

        return dataset.map(
            tokenize,
            remove_columns=dataset.column_names,
            batched=True,
            num_proc=4,  # multiprocessing.cpu_count(),
            load_from_cache_file=False,
        )

最后将所有传入PPOv2Trainer即可开启训练！

    trainer = PPOv2Trainer(
        config=config,
        tokenizer=tokenizer,
        policy=policy,
        ref_policy=ref_policy,
        reward_model=reward_model,
        value_model=value_model,
        train_dataset=prepare_dataset(train_dataset, tokenizer),
        eval_dataset=prepare_dataset(eval_dataset, tokenizer),
    )
    trainer.train()

最后对显存做了一下PPO的实验，使用deepspeed的zero3，可以降低很多：

Reference

1、https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo

作者：淡水，