TRL 文件

快速入門

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

快速入門

它是如何工作的?

透過 PPO 微調語言模型大致包含三個步驟:

  1. 展開(Rollout):語言模型根據一個查詢(query)生成響應或續寫,這個查詢可以是句子的開頭。
  2. 評估(Evaluation):使用一個函式、模型、人類反饋或它們的某種組合來評估查詢和響應。重要的是,這個過程應該為每個查詢/響應對產生一個標量值。最佳化的目標是最大化這個值。
  3. 最佳化(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 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.