Hugging Face 标志

TRL 文档

PPO 训练器

Hugging Face's logo
加入 Hugging Face 社区

并获取增强文档体验

开始使用

PPO 训练器

TRL 支持用于在任意的奖励信号上使用强化学习训练语言模型的 PPO 训练器。奖励信号可以来自手工规则、度量标准或使用奖励模型从偏好数据中获取。有关完整示例，请参阅 examples/notebooks/gpt2-sentiment.ipynb。该训练器深受 OpenAI 的学习总结工作 的启发。

第一步是训练您的 SFT 模型（请参阅 SFTTrainer），以确保我们用于训练的数据对 PPO 算法来说是分布内的。此外，我们还需要训练一个奖励模型（请参阅 RewardTrainer），它将用于使用 PPO 算法优化 SFT 模型。

PPO 如何工作

通过 PPO 调整语言模型大致分为三个步骤

  1. 运行：语言模型根据查询生成响应或延续，这可能是一句话的开始。
  2. 评估：使用函数、模型、人工反馈或它们的组合来评估查询和响应。重要的是，此过程应该为每个查询/响应对产生一个标量值。
  3. 优化：这是最复杂的一部分。在优化步骤中，使用查询/响应对计算序列中标记的 log-概率。这使用训练的和参考模型（通常是微调前的预训练模型）来完成。这两个输出的 KL 散度被用作额外的奖励信号，以确保生成的响应不要偏离参考语言模型太远。然后使用 PPO 训练活动语言模型。

以下图中展示了此过程

图：流程图。

预期的数据集格式

PPOTrainer 预期将生成响应与从奖励模型获得的奖励相匹配的查询对齐。在 PPO 算法的每一步，我们从数据集中采样一批提示，然后使用这些提示从 SFT 模型生成响应。接下来，奖励模型用于计算生成响应的奖励。最后，使用 PPO 算法利用这些奖励来优化 SFT 模型。

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

以下是一个使用HuggingFaceH4/cherry_picked_prompts数据集的示例。

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? "
    ]
}

使用 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

如上所述，奖励可以通过任何返回字符串单个值的函数生成，无论是简单的规则（例如，字符串长度），度量（例如，BLEU），还是基于人类偏好的奖励模型。在示例中，我们使用奖励模型，并使用 transformers.pipeline 进行初始化，以便于使用。

from transformers import pipeline

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

最后，我们使用 tokenizer 对数据集进行预分词，以确保在训练循环中高效地生成响应

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,
    dataset=dataset,
    tokenizer=tokenizer,
)

启动训练循环

由于 PPOTrainer 每一步都需要一个活动的 reward，因此我们需要定义一种方法来在 PPO 算法的每一步获取奖励。在此示例中，我们将使用上面初始化的情感 reward_model。

为了指导生成过程，我们使用 generation_kwargs，这些参数在每一步传递给 model.generate 方法，用于 SFT 模型。更详细的例子可以在 这里 找到。

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


epochs = 10
for epoch in tqdm(range(epochs), "epoch: "):
    for batch in tqdm(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_pretrained("my_ppo_model")

日志记录

在训练和评估时，我们记录以下指标

  • stats：PPO 算法的统计数据，包括损失、熵等。
  • batch：用于训练 SFT 模型的数据批次。
  • rewards：从奖励模型获得的奖励。

PPOTrainer

类 trl.PPOTrainer

< >

( config: 可选 = None model: 可选 = None ref_model: 可选 = None tokenizer: 可选 = None dataset: 合并 = None optimizer: 可选 = None data_collator: 可选 = None num_shared_layers: 可选 = None lr_scheduler: 可选 = None training_data_collator: 可选 = None )

参数

  • **config** (PPOConfig) — 用于PPOTrainer的配置对象。检查PPOConfig的文档了解更多详细信息。
  • **model** (PreTrainedModelWrapper) — 要优化的模型，带有值头的Hugging Face transformer模型。检查PreTrainedModelWrapper的文档了解更多详细信息。
  • **ref_model** (PreTrainedModelWrapper，可选) — 用于KL惩罚的参考模型，带有闲谈语言建模头Hugging Face transformer模型。检查PreTrainedModelWrapper的文档了解更多详细信息。如果没有提供参考模型，训练器将创建一个与待优化的模型具有相同架构的参考模型，并具有共享层。
  • *tokenizer** (PreTrainedTokenizerBase) — 用于编码数据的分词器。请参阅 transformers.PreTrainedTokenizer 和 transformers.PreTrainedTokenizerFast 的文档以获取更多详细信息。
  • *dataset** (Union[torch.utils.data.Dataset, datasets.Dataset], optional) — PyTorch 数据集或 Hugging Face 数据集。用于创建 PyTorch 数据加载器。如果没有提供数据集，则需要在训练器外部创建数据加载器。用户需要设计自己的数据加载器并确保使用的批次大小与配置对象中指定的批次大小相同。
  • *optimizer** (torch.optim.Optimizer, optional) — 用于训练的优化器。如果没有提供优化器，则训练器将创建一个 Adam 优化器，其学习率由配置对象中指定的学习率确定。
  • **data_collator** (DataCollatorForLanguageModeling, optional) — 用于训练的数据编集器，以及传递给数据加载器
  • **num_shared_layers** (int, optional) — 模型和参考模型之间需要共享的层数，如果没有传递参考模型。如果没有提供数值，则所有层都将共享。
  • **lr_scheduler** (torch.optim.lr_scheduler, optional) — 用于训练的学习率调度器。

PPOTrainer 使用近端策略优化来优化语言模型。注意，这个训练器对 OpenAI 的原始学习总结方法有很大启发，这里是学习总结的 GitHub 链接：https://github.com/openai/summarize-from-feedback

批量前向传递

< >

( model: PreTrainedModelWrapper queries: Tensor responses: Tensor model_inputs: dict return_logits: bool = False response_masks: Optional = None ) → (tuple)

参数

  • queries (torch.LongTensor) — 列出包含编码查询的张量，形状为 (batch_size, query_length)
  • responses (torch.LongTensor) — 包含编码响应的张量列表，形状为（《batch_size”，“response_length”）
  • return_logits (bool，可选，默认为False）— 是否返回所有ologits。如果不需要logsits来减少内存消耗，则设置为False。

返回值

(元组)

  • all_logprobs (torch.FloatTensor)：响应的对数概率，形状为（《batch_size”，“response_length”）
  • all_ref_logprobs (torch.FloatTensor)：响应的对数概率，形状为（《batch_size”，“response_length”）
  • all_values (torch.FloatTensor)：响应的值，形状为（《batch_size”，“response_length”）

在多个批次中计算模型输出。

compute_rewards

< >

( scores: FloatTensor logprobs: FloatTensor ref_logprobs: FloatTensor masks: LongTensor ) → torch.FloatTensor

参数

  • scores (torch.FloatTensor) — 奖励模型的分数，形状为 (batch_size)
  • logprobs (torch.FloatTensor) — 模型的对数概率，形状为 (batch_size, response_length)
  • ref_logprobs (torch.FloatTensor) — 参考模型的对数概率，形状为 (batch_size, response_length)

返回值

torch.FloatTensor

每个标记的奖励，形状为 (batch_size, response_length) torch.FloatTensor：非分数奖励，形状为 (batch_size, response_length) torch.FloatTensor：KL惩罚，形状为 (batch_size, response_length)

从分数和KL惩罚中计算每个标记的奖励。

创建模型卡片

< >

( path: str model_name: Optional = 'TRL Model' )

参数

  • path (str) — 保存模型卡片的路径。
  • model_name (str, 可选) — 模型名称，默认为 TRL Model。

创建并保存一个 TRL 模型的模型卡片。

gather_stats

< >

( stats ) → dict[str, Any]

参数

  • stats (dict[str, Any]) —
  • a dictionary of stats to be gathered. The stats should contain torch tensors. —

返回值

dict[str, Any]

收集所有进程的统计状态。

A dictionary of stats with the tensors gathered.

在分布式训练的上下文中非常有用。

< >

( query_tensor: 联合 length_sampler: 可选 = None batch_size: int = 4 return_prompt: bool = True generate_ref_response: bool = False **generation_kwargs ) → torch.LongTensor

参数

  • query_tensor (torch.LongTensor) — 一个形状为（seq_len）的tensor，包含查询token或（seq_len）形状的tensor列表。
  • length_sampler (Callable, 可选) — 返回新生成标记数量的可调用函数。
  • batch_size (int, *可选) — 用于生成的批处理大小，默认值为 4。
  • return_prompt (bool, 可选) — 如果设置为 False，则不会返回提示信息，仅返回新生成的标记，默认为 True。
  • generate_ref_response (bool, 可选) — 如果设置为 True，则也生成参考响应，默认为 False。
  • 生成参数 (dict[str, Any]) — 生成时的关键字参数。

返回值

torch.LongTensor

包含响应令牌的形状为(<code>batch_size</code>, <code>gen_len</code>)的张量。

根据查询张量使用模型生成响应。调用模型的 <code>generate</code> 方法。

日志统计

< >

( 统计数据: dict 批次: dict 奖励: List 要记录的列: Iterable = ('query', 'response') )

参数

  • 统计数据 (dict[str, Any]) — 训练统计数据字典。
  • 批量（dict[str, Any]）—包含查询和响应的批量数据字典。
  • 奖励（code>List[torch.FloatTensor]）—奖励张量。

记录所有训练统计信息的功能。在每个时代的末尾调用它。

损失

< >

( old_logprobs: FloatTensor values: FloatTensor logits: FloatTensor vpreds: FloatTensor logprobs: FloatTensor mask: LongTensor advantages: FloatTensor returns: FloatTensor )

参数

  • old_logprobs (torch.FloatTensor) —— 模型的对数概率，形状为 (batch_size，response_length)
  • values (torch.FloatTensor) —— 值头的值，形状为 (batch_size，response_length)
  • rewards (torch.FloatTensor) —— 奖励模型中的奖励，形状为 (batch_size，response_length)
  • logits (torch.FloatTensor) —— 模型的logits，形状为 (batch_size，response_length，vocab_size)
  • v_pred (torch.FloatTensor) — 值头部的值，形状(batch_size, response_length)
  • logprobs (torch.FloatTensor) — 模型的对数概率，形状(batch_size, response_length)

计算策略和值损耗。

prepare_dataloader

< >

( dataset: Union data_collator = None ) → torch.utils.data.DataLoader

参数

  • dataset (Union[torch.utils.data.Dataset, datasets.Dataset]) — PyTorch 数据集或 Hugging Face 数据集。如果传入 Hugging Face 数据集，则数据集将被预处理，删除模型未使用的列。
  • data_collator (Optional[function]) — 数据收集函数。

返回值

torch.utils.data.DataLoader

PyTorch 数据加载器

为训练准备数据加载器。

record_step_stats

< >

( kl_coef: float **data ) → stats (dict)

参数

  • kl_coef (float) — KL系数
  • data (dict) — 训练步骤数据字典

返回值

stats (dict)

训练步骤统计字典

记录训练步骤统计。

step

< >

( queries: 列表 responses: 列表 scores: 列表 response_masks: 可选 = None ) → dict[str, Any]

参数

  • queries (Listtorch.LongTensor) — 包含编码查询的张量列表，形状为(<code>query\_length</code>)
  • responses (Listtorch.LongTensor) — 包含编码响应的张量列表，形状为(<code>response\_length</code>)
  • scores (Listtorch.FloatTensor) — 包含分数的张量列表。
  • response_masks (Listtorch.FloatTensor, 可选)) — 包含响应标记遮罩的张量列表。

返回值

dict[str, Any]

训练统计概要

给定查询列表、模型响应和奖励运行 PPO 最优化步骤。

train\_minibatch

< >

( old_logprobs: FloatTensor values: FloatTensor logprobs: FloatTensor logits: FloatTensor vpreds: FloatTensor mask: LongTensor advantages: FloatTensor returns: FloatTensor ) → train_stats (dict[str, torch.Tensor])

参数

  • logprobs (torch.FloatTensor) — 模型的对数概率，形状 [mini_batch_size, response_length]
  • values (torch.FloatTensor) — 值头部值，形状 [mini_batch_size, response_length]
  • query (torch.LongTensor) — 编码查询，形状 [mini_batch_size, query_length]
  • response (torch.LongTensor) — 编码响应，形状 [mini_batch_size, response_length]
  • model_input (torch.LongTensor) — 连接查询和响应，形状 [mini_batch_size, query_length+response_length]

返回值

train_stats (dict[str, torch.Tensor])

训练统计字典

训练一个PPO小批量

类 trl.PPOConfig

< >

( exp_name: str = 'doc-buil' seed: int = 0 log_with: Optional = None task_name: Optional = None model_name: Optional = 'gpt2' query_dataset: Optional = 'imdb' reward_model: Optional = 'sentiment-analysis:lvwerra/distilbert-imdb' remove_unused_columns: bool = True tracker_kwargs: Annotated = <factory> accelerator_kwargs: Annotated = <factory> project_kwargs: Annotated = <factory> tracker_project_name: str = 'trl' push_to_hub_if_best_kwargs: Annotated = <factory> steps: int = 20000 learning_rate: float = 1.41e-05 adap_kl_ctrl: bool = True init_kl_coef: Optional = 0.2 kl_penalty: Literal = 'kl' target: Optional = 6 horizon: Optional = 10000 gamma: float = 1 lam: float = 0.95 cliprange: float = 0.2 cliprange_value: float = 0.2 vf_coef: float = 0.1 batch_size: int = 128 forward_batch_size: Optional = None mini_batch_size: int = 128 gradient_accumulation_steps: int = 1 world_size: Annotated = None ppo_epochs: int = 4 max_grad_norm: Optional = None optimize_cuda_cache: Optional = None optimize_device_cache: Optional = False early_stopping: bool = False target_kl: float = 1 compare_steps: int = 1 ratio_threshold: float = 10.0 use_score_scaling: bool = False use_score_norm: bool = False score_clip: Optional = None whiten_rewards: bool = False gradient_checkpointing: bool = False is_encoder_decoder: Optional = None is_peft_model: Optional = None backward_batch_size: Annotated = None global_backward_batch_size: Annotated = None global_batch_size: Annotated = None )

PPOTrainer的配置类

< > GitHub更新

←监督细微调整 PPOv2 训练器→