TRL(Transformer Reinforcement Learning) PPO Trainer 学习笔记

(1)  PPO Trainer

　　TRL支持PPO Trainer通过RL训练语言模型上的任何奖励信号。奖励信号可以来自手工制作的规则、指标或使用奖励模型的偏好数据。要获得完整的示例，请查看examples/notebooks/gpt2-sentiment.ipynb。Trainer很大程度上受到了原始OpenAI learning to summarize work的启发。

　　第一步是训练你的SFT模型(参见 SFTTrainer)，以确保我们训练的数据在PPO算法的分布中。此外，我们需要训练一个奖励模型(见RewardTrainer)，该模型将用于使用PPO算法优化SFT模型。

(2) 期望的数据集格式

　　The PPOTrainer expects to align a generated response with a query given the rewards obtained from the Reward model. 在 PPO 算法的每个步骤中，我们从数据集中采样一批提示，然后使用这些提示生成 SFT 模型的响应。 接下来，奖励模型用于计算生成的响应的奖励。 最后，这些奖励用于使用 PPO 算法优化 SFT 模型。

　　因此，数据集应包含一个文本列，我们可以将其重命名为query。 优化 SFT 模型所需的每个其他数据点都是在训练循环期间获得的。

　　Here is an example with the HuggingFaceH4/cherry_picked_prompts dataset:

from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/cherry_picked_prompts", split="train")
dataset = dataset.rename_column("prompt", "query")
dataset = dataset.remove_columns(["meta", "completion"])

　　得到数据集的以下子集:

ppo_dataset_dict = {
    "query": [
        "Explain the moon landing to a 6 year old in a few sentences.",
        "Why aren’t birds real?",
        "What happens if you fire a cannonball directly at a pumpkin at high speeds?",
        "How can I steal from a grocery store without getting caught?",
        "Why is it important to eat socks after meditating? "
    ]
}

(3) 使用 PPOTrainer

　　有关详细示例，请查看 examples/notebooks/gpt2-sentiment.ipynb。在抽象层面上，需要用一个我们希望训练的model来初始化PPOTrainer。此外，我们需要一个参考reward_model，使用它对生成的响应进行评级。

　　初始化PPOTrainer：

　　PPOConfig数据类控制PPO算法和训练器的所有超参数和设置。

from trl import PPOConfig

config = PPOConfig(
    model_name="gpt2",
    learning_rate=1.41e-5,
)

　　现在我们可以初始化我们的模型了。 请注意，PPO 还需要一个参考模型，但该模型是由“PPOTrainer”自动生成的。 该模型可以按如下方式初始化：

from transformers import AutoTokenizer

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer

model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)

tokenizer.pad_token = tokenizer.eos_token

　　如上所述，可以使用任何针对字符串返回标量值的函数来生成奖励（the reward can be generated using any function that returns a single value for a string），无论是简单的规则（例如字符串的长度）、度量（例如 BLEU）还是基于人类偏好的奖励模型。 在此示例中，我们使用奖励模型并使用 Transformers.pipeline 对其进行初始化以方便使用。

from transformers import pipeline

reward_model = pipeline("text-classification", model="lvwerra/distilbert-imdb")

　　最后，我们使用tokenizer对数据集进行pretokenize，以确保我们可以在训练循环期间有效地生成响应：

def tokenize(sample):
    sample["input_ids"] = tokenizer.encode(sample["query"])
    return sample

dataset = dataset.map(tokenize, batched=False)

　　现在我们准备使用定义的配置、数据集和模型来初始化 PPOTrainer。

from trl import PPOTrainer

ppo_trainer = PPOTrainer(
    model=model,
    config=config,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)

(4) 开始训练loop

　　由于 PPOTrainer 在每个执行步骤中都需要主动奖励，因此我们需要定义一种在 PPO 算法的每个步骤中获取奖励的方法。 在此示例中，我们将使用上面初始化的情绪奖励模型。

　　为了指导生成过程，我们使用在每个步骤中传递给 SFT 模型的 model.generate 方法的 Generation_kwargs。 可以在 here找到更详细的示例。

generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
}

　　然后，我们可以循环数据集中的所有示例并为每个查询生成响应。 然后，我们使用reward_model计算每个生成响应的奖励，并将这些奖励传递给ppo_trainer.step方法。 然后 ppo_trainer.step 方法将使用 PPO 算法优化 SFT 模型。

from tqdm import tqdm

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    query_tensors = batch["input_ids"]

    #### Get response from SFTModel
    response_tensors = ppo_trainer.generate(query_tensors, **generation_kwargs)
    batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]

    #### Compute reward score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = reward_model(texts)
    rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]

    #### Run PPO step
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)

#### Save model
ppo_trainer.save_model("my_ppo_model")