快速入门

它是如何工作的？

通过 PPO 微调语言模型大致包含三个步骤：

展开（Rollout）：语言模型根据一个查询（query）生成响应或续写，这个查询可以是句子的开头。
评估（Evaluation）：使用一个函数、模型、人类反馈或它们的某种组合来评估查询和响应。重要的是，这个过程应该为每个查询/响应对产生一个标量值。优化的目标是最大化这个值。
优化（Optimization）：这是最复杂的部分。在优化步骤中，查询/响应对被用来计算序列中词元（token）的对数概率。这是通过正在训练的模型和一个参考模型来完成的，参考模型通常是微调前的预训练模型。两个输出之间的 KL 散度被用作一个额外的奖励信号，以确保生成的响应不会偏离参考语言模型太远。然后，活动语言模型会使用 PPO 进行训练。

整个过程如下图所示：

最小示例

以下代码展示了上述步骤。

# 0. imports
import torch
from transformers import GPT2Tokenizer

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer


# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained("gpt2")
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# 2. initialize trainer
ppo_config = {"mini_batch_size": 1, "batch_size": 1}
config = PPOConfig(**ppo_config)
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)

# 3. encode a query
query_txt = "This morning I went to the "
query_tensor = tokenizer.encode(query_txt, return_tensors="pt").to(model.pretrained_model.device)

# 4. generate model response
generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "max_new_tokens": 20,
}
response_tensor = ppo_trainer.generate([item for item in query_tensor], return_prompt=False, **generation_kwargs)
response_txt = tokenizer.decode(response_tensor[0])

# 5. define a reward for response
# (this could be any reward such as human feedback or output from another model)
reward = [torch.tensor(1.0, device=model.pretrained_model.device)]

# 6. train model with ppo
train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)

通常，你会在一个 for 循环中运行步骤 3-6，并对许多不同的查询运行它。你可以在示例部分找到更实际的例子。

如何使用训练好的模型

在训练一个 AutoModelForCausalLMWithValueHead 后，你可以直接在 transformers 中使用该模型。


# .. Let's assume we have a trained model using `PPOTrainer` and `AutoModelForCausalLMWithValueHead`

# push the model on the Hub
model.push_to_hub("my-fine-tuned-model-ppo")

# or save it locally
model.save_pretrained("my-fine-tuned-model-ppo")

# load the model from the Hub
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("my-fine-tuned-model-ppo")

如果你想使用价值头（value head），例如继续训练，你也可以用 AutoModelForCausalLMWithValueHead 加载你的模型。

from trl.model import AutoModelForCausalLMWithValueHead

model = AutoModelForCausalLMWithValueHead.from_pretrained("my-fine-tuned-model-ppo")

< > 在 GitHub 上更新

TRL

快速入门

它是如何工作的？

最小示例

如何使用训练好的模型