RLOO 训练器

TRL 支持使用 REINFORCE 留一法 (RLOO) 训练 LLM。其思想是，RLOO 不使用价值函数，而是为每个提示生成 K 个补全。对于每个补全，RLOO 使用其他 K-1 个补全的平均分数作为基线来计算优势。RLOO 还将整个补全建模为单个动作，而 PPO 将每个 token 建模为一个动作。请注意，REINFORCE / A2C 是 PPO 的一个特例，当 PPO epoch 数为 1 且 mini-batch 数为 1 时，这正是我们在 TRL 中实现 RLOO 的方式。

参考文献

开始使用

要运行 RLOO 脚本以确保训练器可以运行，您可以运行以下命令来训练带有虚拟奖励模型的 RLOO 模型。

python examples/scripts/rloo/rloo.py \
    --dataset_name trl-internal-testing/descriptiveness-sentiment-trl-style \
    --dataset_train_split descriptiveness \
    --learning_rate 3e-6 \
    --output_dir models/minimal/rloo \
    --per_device_train_batch_size 64 \
    --gradient_accumulation_steps 1 \
    --total_episodes 10000 \
    --model_name_or_path EleutherAI/pythia-14m \
    --reward_model_path EleutherAI/pythia-14m \
    --missing_eos_penalty 1.0

记录指标的解释

记录的指标如下。这是一个示例在 Weights and Biases 跟踪的运行

eps: 跟踪每秒的 episode 数。
objective/kl: 当前策略和参考策略之间的平均 Kullback-Leibler (KL) 散度。
objective/entropy: 策略的平均熵，指示策略选择的动作的随机性。
objective/non_score_reward: 来自非分数相关来源的平均奖励，基本上是 beta * kl.sum(1)，其中 beta 是 KL 惩罚系数，kl 是每个 token 的 KL 散度。
objective/rlhf_reward: 平均 RLHF 奖励，即 score - non_score_reward。
objective/scores: 奖励模型/环境返回的平均分数。
policy/approxkl_avg: 连续 PPO 策略之间的平均近似 KL 散度。请注意，这与 objective/kl 不同。
policy/clipfrac_avg: 策略更新的平均裁剪比例，指示策略更新被约束以防止大幅更改的频率。
loss/policy_avg: 平均策略损失，指示策略的性能。
val/clipfrac_avg: 价值函数更新的平均裁剪比例，类似于 policy/clipfrac_avg，但用于价值函数。
policy/entropy_avg: 训练期间策略的平均熵，指示策略动作的多样性。
val/ratio: 当前策略概率与旧策略概率的平均比率，提供策略更改程度的度量。
val/ratio_var: val/ratio 的方差，指示策略更改的可变性。
val/num_eos_tokens: 生成的序列结束 (EOS) token 的数量，这可以指示完整响应的数量。
lr: lr：优化器使用的当前学习率。
episode: episode：训练过程中的当前全局步骤或 episode 计数。

实用技巧

调试技巧：objective/rlhf_reward：这是 RLHF 训练的最终目标。如果训练按预期进行，则此指标应持续上升。
调试技巧：val/ratio：这个数字应该在 1.0 左右浮动，并且通过 PPO 代理损失的 --cliprange 0.2 进行裁剪。因此，如果这个 ratio 太高（如 2.0 或 1000.0）或太小（如 0.1），则意味着连续策略之间的更新过于剧烈。您应该尝试理解为什么会发生这种情况并尝试修复它。
内存技巧：如果您的内存不足，您可以尝试减少 --per_device_train_batch_size 或增加 --gradient_accumulation_steps 以减少内存占用。
内存技巧：如果您有多个 GPU，您还可以使用 DeepSpeed stage 3 运行训练以减少内存占用 accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml。
使用技巧：我们建议通过 --missing_eos_penalty 使用“EOS 技巧”，这会从不以 EOS token 结尾的补全的分数中减去一个静态标量惩罚。这可以帮助模型学习生成更连贯的补全。

我的模型到底在做什么？

为了帮助您了解您的模型正在做什么，我们会定期记录模型的一些示例补全。这是一个补全的示例。在示例在 Weights and Biases 跟踪的运行中，它看起来像下面这样，让您可以在训练的不同阶段看到模型的响应。默认情况下，我们在训练期间生成 --num_sample_generations 10，但您可以自定义生成数量。

在日志中，采样的生成结果如下所示

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │
│                                 │                                 │          │
│ POST: Hi,                       │                                 │          │
│ I'm 22, and I have been with my │                                 │          │
│ girlfriend for 5 years now. We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.                      │                                 │          │
│                                 │                                 │          │
│ Problem, I recently started to  │                                 │          │
│ have feelings for an other      │                                 │          │
│ person (a friend). This person  │                                 │          │
│ has had a boyfriend for now 3   │                                 │          │
│ years, and has absolutely no    │                                 │          │
│ ideas. Those feelings were so   │                                 │          │
│ strong, it was hard to hide     │                                 │          │
│ them. After 2 months of me      │                                 │          │
│ being distant and really sad,   │                                 │          │
│ my girlfriend forced me to say  │                                 │          │
│ what was bothering me. I'm not  │                                 │          │
│ a good liar, and now she knows. │                                 │          │
│                                 │                                 │          │
│ We decided to give us a week    │                                 │          │
│ alone, I went to my parents.    │                                 │          │
│                                 │                                 │          │
│ Now, I'm completely lost. I     │                                 │          │
│ keep on thinking about this     │                                 │          │
│ person, and I hate that. I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.                    │                                 │          │
│                                 │                                 │          │
│ What do I do? It's been 3       │                                 │          │
│ months now, and I'm just        │                                 │          │
│ desperate.                      │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
├─────────────────────────────────┼─────────────────────────────────┼──────────┤
│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │
│                                 │ TV. I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down. She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality. I had no  │                                 │          │
│ lock on my door; but I had a    │                                 │          │
│ dresser right next to it, so I  │                                 │          │
│ pulled one of the drawers out   │                                 │          │
│ enough so that it caused the    │                                 │          │
│ door to not be openable. Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘

实现细节

RLOOTrainer 的大部分基于 PPO 实现，该实现基于使用 PPO 的 RLHF 的 N+ 个实现细节：关于 TL;DR 摘要的案例研究。

下面是 RLOO 的向量化优势计算

def test_rloo_reward():
    local_batch_size = 3
    rloo_k = 4
    rlhf_reward = torch.tensor([
        1, 2, 3, # first rlhf reward for three prompts
        2, 3, 4, # second rlhf reward for three prompts
        5, 6, 7, # third rlhf reward for three prompts
        8, 9, 10, # fourth rlhf reward for three prompts
    ]).float() # here we have 3 prompts which have 4 completions each

    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
    advantages = torch.zeros_like(rlhf_reward)
    for i in range(0, len(advantages), local_batch_size):
        other_response_rlhf_rewards = []
        for j in range(0, len(advantages), local_batch_size):
            if i != j:
                other_response_rlhf_rewards.append(rlhf_reward[j : j + local_batch_size])
        advantages[i : i + local_batch_size] = rlhf_reward[i : i + local_batch_size] - torch.stack(other_response_rlhf_rewards).mean(0)
    
    assert (1 - (2 + 5 + 8) / 3 - advantages[0].item()) < 1e-6  # First rlhf reward for the first prompt
    assert (6 - (3 + 2 + 9) / 3 - advantages[7].item()) < 1e-6  # Third rlhf reward for the second prompt

    # Vectorized implementation
    rlhf_reward = rlhf_reward.reshape(rloo_k, local_batch_size)
    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
    vec_advantages = rlhf_reward - baseline
    torch.testing.assert_close(vec_advantages.flatten(), advantages)

基准实验

为了验证 RLOO 实现是否有效，我们在 1B 模型上运行了实验。以下是我们用来运行实验的命令。我们直接从使用 PPO 的 RLHF 的 N+ 个实现细节：关于 TL;DR 摘要的案例研究中获取 SFT/RM 模型。

accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
    --output_dir models/minimal/rloo_tldr \
    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style \
    --dataset_test_split validation \
    --num_ppo_epochs 2 \
    --num_mini_batches 2 \
    --learning_rate 3e-6 \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 16 \
    --total_episodes 1000000 \
    --model_name_or_path EleutherAI/pythia-1b-deduped \
    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
    --local_rollout_forward_batch_size 16 \
    --missing_eos_penalty 1.0 \
    --stop_token eos \
    --kl_coef 0.03

检查点和实验跟踪可在以下位置获得

为了评估，我们使用 vLLM 加载检查点，并使用 GPT-4o mini 作为评判模型来评估生成的 TL;DR 与参考 TL;DR 的对比。有关如何使用评判器的更多信息，请参阅评判器。

$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 51.20%

RLOO 检查点获得了 51.2% 的首选率，而 SFT 检查点的首选率为 33.0%。这是一个很好的迹象，表明 RLOO 训练正在按预期工作。

指标

# pip install openrlbenchmark==0.2.1a5
# see https://github.com/openrlbenchmark/openrlbenchmark#get-started for documentation
# to use it, change `?we=huggingface&wpn=trl` to your own project and `?tag=pr-1540` to your own tag
python -m openrlbenchmark.rlops_multi_metrics \
    --filters '?we=huggingface&wpn=trl&xaxis=train/episode&ceik=output_dir&cen=sft_model_path&metrics=train/objective/rlhf_reward&metrics=train/objective/scores&metrics=train/objective/kl&metrics=train/objective/non_score_reward&metrics=train/objective/entropy&metrics=train/policy/approxkl_avg&metrics=train/policy/clipfrac_avg&metrics=train/loss/policy_avg&metrics=train/policy/entropy_avg&metrics=train/val/ratio&metrics=train/val/ratio_var&metrics=train/val/num_eos_tokens&metrics=train/lr&metrics=train/eps' \
        "cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr?tag=pr-1540" \
    --env-ids models/minimal/rloo_tldr \
    --pc.ncols 4 \
    --pc.ncols-legend 1 \
    --pc.xlabel "Episode" \
    --output-filename benchmark/trl/pr-1540/rloo \
    --scan-history

Reinforce++

Jian Hu 的 Reinforce++ 报告提出了几个优化技巧，以增强 RLHF 的性能和稳定性。它们包括

裁剪奖励：将奖励值限制在特定范围内，以减轻极端奖励对模型更新的影响，从而防止梯度爆炸
归一化奖励：将奖励缩放为平均值为 0，标准差为 1，这有助于稳定训练过程
归一化优势：将优势缩放为平均值为 0，标准差为 1，这有助于稳定训练过程
使用 token 级别 KL 惩罚（定义为报告的公式 (1)）而不是序列级别 KL 惩罚（默认）

这些选项可通过 RLOOConfig 类中的相应参数获得。

RLOOTrainer

class trl.RLOOTrainer

< source >

( config: RLOOConfig processing_class: typing.Union[transformers.tokenization_utils_base.PreTrainedTokenizerBase, transformers.image_processing_utils.BaseImageProcessor, transformers.feature_extraction_utils.FeatureExtractionMixin, transformers.processing_utils.ProcessorMixin, NoneType] policy: Module ref_policy: Module reward_model: typing.Union[torch.nn.modules.module.Module, typing.Callable[[list[str]], list[float]]] train_dataset: Dataset data_collator: typing.Optional[transformers.data.data_collator.DataCollatorWithPadding] = None eval_dataset: typing.Union[datasets.arrow_dataset.Dataset, dict[str, datasets.arrow_dataset.Dataset], NoneType] = None optimizers: tuple = (None, None) callbacks: typing.Optional[list[transformers.trainer_callback.TrainerCallback]] = None )

create_model_card

< source >

( model_name: typing.Optional[str] = None dataset_name: typing.Optional[str] = None tags: typing.Union[str, list[str], NoneType] = None )

参数

model_name (str 或 None, 可选, 默认为 None) — 模型的名称。
dataset_name (str 或 None, 可选, 默认为 None) — 用于训练的数据集名称。
tags (str, list[str] 或 None, 可选, 默认为 None) — 与模型卡关联的标签。

使用 Trainer 可用的信息创建模型卡的草稿。

RLOOConfig

class trl.RLOOConfig

< source >

( output_dir: typing.Optional[str] = None overwrite_output_dir: bool = False do_train: bool = False do_eval: bool = False do_predict: bool = False eval_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'no' prediction_loss_only: bool = False per_device_train_batch_size: int = 8 per_device_eval_batch_size: int = 8 per_gpu_train_batch_size: typing.Optional[int] = None per_gpu_eval_batch_size: typing.Optional[int] = None gradient_accumulation_steps: int = 1 eval_accumulation_steps: typing.Optional[int] = None eval_delay: typing.Optional[float] = 0 torch_empty_cache_steps: typing.Optional[int] = None learning_rate: float = 5e-05 weight_decay: float = 0.0 adam_beta1: float = 0.9 adam_beta2: float = 0.999 adam_epsilon: float = 1e-08 max_grad_norm: float = 1.0 num_train_epochs: float = 3.0 max_steps: int = -1 lr_scheduler_type: typing.Union[transformers.trainer_utils.SchedulerType, str] = 'linear' lr_scheduler_kwargs: typing.Union[dict, str, NoneType] = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: typing.Optional[str] = 'passive' log_level_replica: typing.Optional[str] = 'warning' log_on_each_node: bool = True logging_dir: typing.Optional[str] = None logging_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps' logging_first_step: bool = False logging_steps: float = 500 logging_nan_inf_filter: bool = True save_strategy: typing.Union[transformers.trainer_utils.SaveStrategy, str] = 'steps' save_steps: float = 500 save_total_limit: typing.Optional[int] = None save_safetensors: typing.Optional[bool] = True save_on_each_node: bool = False save_only_model: bool = False restore_callback_states_from_checkpoint: bool = False no_cuda: bool = False use_cpu: bool = False use_mps_device: bool = False seed: int = 42 data_seed: typing.Optional[int] = None jit_mode_eval: bool = False use_ipex: bool = False bf16: bool = False fp16: bool = False fp16_opt_level: str = 'O1' half_precision_backend: str = 'auto' bf16_full_eval: bool = False fp16_full_eval: bool = False tf32: typing.Optional[bool] = None local_rank: int = -1 ddp_backend: typing.Optional[str] = None tpu_num_cores: typing.Optional[int] = None tpu_metrics_debug: bool = False debug: typing.Union[str, list[transformers.debug_utils.DebugOption]] = '' dataloader_drop_last: bool = False eval_steps: typing.Optional[float] = None dataloader_num_workers: int = 0 dataloader_prefetch_factor: typing.Optional[int] = None past_index: int = -1 run_name: typing.Optional[str] = None disable_tqdm: typing.Optional[bool] = None remove_unused_columns: typing.Optional[bool] = True label_names: typing.Optional[list[str]] = None load_best_model_at_end: typing.Optional[bool] = False metric_for_best_model: typing.Optional[str] = None greater_is_better: typing.Optional[bool] = None ignore_data_skip: bool = False fsdp: typing.Union[list[transformers.trainer_utils.FSDPOption], str, NoneType] = '' fsdp_min_num_params: int = 0 fsdp_config: typing.Union[dict, str, NoneType] = None tp_size: typing.Optional[int] = 0 fsdp_transformer_layer_cls_to_wrap: typing.Optional[str] = None accelerator_config: typing.Union[dict, str, NoneType] = None deepspeed: typing.Union[dict, str, NoneType] = None label_smoothing_factor: float = 0.0 optim: typing.Union[transformers.training_args.OptimizerNames, str] = 'adamw_torch' optim_args: typing.Optional[str] = None adafactor: bool = False group_by_length: bool = False length_column_name: typing.Optional[str] = 'length' report_to: typing.Union[NoneType, str, list[str]] = None ddp_find_unused_parameters: typing.Optional[bool] = None ddp_bucket_cap_mb: typing.Optional[int] = None ddp_broadcast_buffers: typing.Optional[bool] = None dataloader_pin_memory: bool = True dataloader_persistent_workers: bool = False skip_memory_metrics: bool = True use_legacy_prediction_loop: bool = False push_to_hub: bool = False resume_from_checkpoint: typing.Optional[str] = None hub_model_id: typing.Optional[str] = None hub_strategy: typing.Union[transformers.trainer_utils.HubStrategy, str] = 'every_save' hub_token: typing.Optional[str] = None hub_private_repo: typing.Optional[bool] = None hub_always_push: bool = False gradient_checkpointing: bool = False gradient_checkpointing_kwargs: typing.Union[dict, str, NoneType] = None include_inputs_for_metrics: bool = False include_for_metrics: list = <factory> eval_do_concat_batches: bool = True fp16_backend: str = 'auto' push_to_hub_model_id: typing.Optional[str] = None push_to_hub_organization: typing.Optional[str] = None push_to_hub_token: typing.Optional[str] = None mp_parameters: str = '' auto_find_batch_size: bool = False full_determinism: bool = False torchdynamo: typing.Optional[str] = None ray_scope: typing.Optional[str] = 'last' ddp_timeout: typing.Optional[int] = 1800 torch_compile: bool = False torch_compile_backend: typing.Optional[str] = None torch_compile_mode: typing.Optional[str] = None include_tokens_per_second: typing.Optional[bool] = False include_num_input_tokens_seen: typing.Optional[bool] = False neftune_noise_alpha: typing.Optional[float] = None optim_target_modules: typing.Union[NoneType, str, list[str]] = None batch_eval_metrics: bool = False eval_on_start: bool = False use_liger_kernel: typing.Optional[bool] = False eval_use_gather_object: typing.Optional[bool] = False average_tokens_across_devices: typing.Optional[bool] = False dataset_num_proc: typing.Optional[int] = None num_mini_batches: int = 1 total_episodes: typing.Optional[int] = None local_rollout_forward_batch_size: int = 64 num_sample_generations: int = 10 response_length: int = 53 stop_token: typing.Optional[typing.Literal['eos']] = None stop_token_id: typing.Optional[int] = None temperature: float = 0.7 missing_eos_penalty: typing.Optional[float] = None sft_model_path: str = 'EleutherAI/pythia-160m' world_size: typing.Optional[int] = None num_total_batches: typing.Optional[int] = None micro_batch_size: typing.Optional[int] = None local_batch_size: typing.Optional[int] = None batch_size: typing.Optional[int] = None local_mini_batch_size: typing.Optional[int] = None mini_batch_size: typing.Optional[int] = None exp_name: str = 'rloo_config' reward_model_path: str = 'EleutherAI/pythia-160m' num_ppo_epochs: int = 4 whiten_rewards: bool = False kl_coef: float = 0.05 cliprange: float = 0.2 rloo_k: int = 2 normalize_reward: bool = False reward_clip_range: float = 10.0 normalize_advantage: bool = False token_level_kl: bool = False ds3_gather_for_generation: bool = True )

参数

exp_name (str, 可选, 默认为 os.path.basename(__file__)[ -- -len(".py")]): 此实验的名称。
reward_model_path (str, 可选, 默认为 "EleutherAI/pythia-160m") — 奖励模型的路径。
num_ppo_epochs (int, 可选, 默认为 4) — 训练的 epoch 数量。
whiten_rewards (bool, 可选, 默认为 False) — 是否对奖励进行白化处理。
kl_coef (float, 可选, 默认为 0.05) — KL 系数。
cliprange (float, 可选, 默认为 0.2) — 裁剪范围。
rloo_k (int, 可选, 默认为 2) — REINFORCE 留一法 (RLOO) 每个提示的在线样本数。
normalize_reward (bool, 可选, 默认为 False) — 是否标准化奖励。
reward_clip_range (float, 可选, 默认为 10.0) — 奖励的裁剪范围。
normalize_advantage (bool, 可选, 默认为 False) — 是否标准化优势函数。
token_level_kl (bool, 可选, 默认为 True) — 是否使用 token 级别 KL 惩罚或序列级别 KL 惩罚。
ds3_gather_for_generation (bool, 可选, 默认为 True) — 此设置应用于 DeepSpeed ZeRO-3。如果启用，策略模型权重将在生成时收集，从而提高生成速度。但是，禁用此选项允许训练超出单个 GPU VRAM 容量的模型，但会以较慢的生成速度为代价。

用于 RLOOTrainer 的配置类。

使用 HfArgumentParser，我们可以将此类转换为可以在命令行中指定的 argparse 参数。

< > 在 GitHub 上更新