快速入门
它是如何工作的?
通过 PPO 微调语言模型大致包括三个步骤
- 展开:语言模型根据查询生成响应或延续,查询可能是句子的开头。
- 评估:查询和响应使用函数、模型、人工反馈或它们的组合进行评估。重要的是,此过程应为每个查询/响应对生成标量值。优化旨在最大化此值。
- 优化:这是最复杂的部分。在优化步骤中,使用查询/响应对来计算序列中标记的对数概率。这是使用训练的模型和参考模型完成的,参考模型通常是微调前的预训练模型。两个输出之间的 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")
如果您想使用值头(例如继续训练),也可以使用 `AutoModelForCausalLMWithValueHead` 加载您的模型。
from trl.model import AutoModelForCausalLMWithValueHead
model = AutoModelForCausalLMWithValueHead.from_pretrained("my-fine-tuned-model-ppo")