A2PO

TRL 支持 A*-PO（最优优势回归），如 Kianté Brantley、Mingyu Chen、Zhaolin Gao、Jason D. Lee、Wen Sun、Wenhao Zhan 和 Xuezhou Zhang 的论文 Accelerating RL for LLM Reasoning with Optimal Advantage Regression 中所述。

该论文的摘要如下：

强化学习（RL）已成为微调大型语言模型（LLM）以提升复杂推理能力的强大工具。然而，最先进的策略优化方法通常面临高计算开销和内存消耗的问题，这主要源于每个提示需要多次生成，以及对当前策略的评论家网络或优势估计的依赖。在本文中，我们提出了 A*-PO，这是一种新颖的两阶段策略优化框架，它直接近似最优优势函数，并支持针对推理任务高效训练 LLM。在第一阶段，我们利用来自参考策略的离线采样来估计最优价值函数 V*，从而消除了对昂贵的在线价值估计的需求。在第二阶段，我们使用简单的最小二乘回归损失进行在策略更新，每个提示仅需单次生成。在理论上，我们建立了性能保证，并证明了 KL 正则化的 RL 目标可以在无需复杂探索策略的情况下进行优化。在实证上，A*-PO 在广泛的数学推理基准测试中取得了具有竞争力的性能，与 PPO、GRPO 和 REBEL 相比，训练时间最多减少了 2 倍，峰值内存使用率降低了 30% 以上。

用法

A*-PO 假设二元且可验证的奖励（r ∈ {0, 1}），并分两个阶段运行：

离线价值估计。在训练之前，针对每个提示词，从参考策略中采样num_value_samples个补全，并使用reward_funcs进行评分。随后，估计并缓存每个提示词对应的最优价值V*(x) = β₁·log(mean_i exp(r(x, yᵢ)/β₁))。
在策略回归。在训练过程中，针对每个提示词，从当前策略中生成一个补全。损失函数是隐式奖励 β₂·log(π(y|x)/π_ref(y|x)) 与最优优势 r(x, y) − V*(x) 之间的平方误差。

from trl.experimental.a2po import A2POConfig, A2POTrainer

# A*-PO assumes a binary, verifiable reward in {0, 1}.
def reward_correct(completions, ground_truth, **kwargs):
    return [float(completion.strip() == truth) for completion, truth in zip(completions, ground_truth)]

training_args = A2POConfig(
    output_dir="Qwen2.5-0.5B-A2PO",
    num_value_samples=8,  # Stage 1: samples per prompt from the reference policy to estimate V*
    beta1=0.5,  # Stage 1: KL temperature for the V* estimate
    beta2=1e-3,  # Stage 2: KL temperature for the regression target
)
trainer = A2POTrainer(
    model="Qwen/Qwen2.5-0.5B",
    reward_funcs=reward_correct,
    args=training_args,
    train_dataset=...,
)
trainer.train()

由于 V* 完全基于参考策略的样本进行估计，A*-PO 无法超越参考策略的 Pass@K。官方实现可在 ZhaolinGao/A-PO 中找到。

A2POTrainer

类 trl.experimental.a2po.A2POTrainer

< source >

( model: transformers.modeling_utils.PreTrainedModel | strreward_funcs: collections.abc.Callable[..., list[float]] | list[collections.abc.Callable[..., list[float]]]args: trl.experimental.a2po.a2po_config.A2POConfig | None = Nonetrain_dataset = Noneeval_dataset = Noneprocessing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | None = Nonecallbacks = Noneoptimizers = (None, None) )

参数

model (PreTrainedModel 或 str) — 要训练的模型，或传递给 from_pretrained 的模型标识符（字符串）。
reward_funcs (Callable 或 list[Callable]) — 奖励函数。每个函数接收 prompts 和 completions（以及作为关键字参数传入的数据集列），并返回一个浮点数奖励列表。当提供多个函数时，它们的加权和（参见 A2POConfig.reward_weights）即为标量奖励 r，A*-PO 假设其为二值（位于 {0, 1} 中）。
args (experimental.a2po.A2POConfig，可选) — 此训练器的配置。如果为 None，则使用默认配置。
train_dataset (Dataset，可选) — 训练数据集。必须包含一个 "prompt" 列。
eval_dataset (Dataset，可选) — 评估数据集。
processing_class (PreTrainedTokenizerBase, 可选) — 用于处理数据的处理类。如果为 None，则通过 from_pretrained 从模型名称加载。
callbacks (list[~transformers.TrainerCallback]，可选) — 用于自定义训练循环的回调列表。
optimizers (tuple[~torch.optim.Optimizer, ~torch.optim.lr_scheduler.LambdaLR]，可选，默认为 (None, None)) — 包含优化器和学习率调度器的元组。

Trainer for the A*-PO (Optimal Advantage Regression) method, introduced in Accelerating RL for LLM Reasoning with Optimal Advantage Regression.

A*-PO 分两个阶段运行：

离线价值估计。在训练之前，会为每个训练提示词从参考策略中采样num_value_samples个补全，并使用reward_funcs进行评分。最优价值被估计为V*(x) = beta1 * log(mean_i exp(r(x, y_i) / beta1))，并按提示词进行缓存。
在策略回归。在训练过程中，针对每个提示词，从当前策略中生成一个补全。损失函数是隐式奖励 beta2 * log(pi(y|x) / pi_ref(y|x)) 与最优优势估计 r(x, y) - V*(x) 之间的平方误差。

训练

< source >

( *args**kwargs )

save_model

< source >

( output_dir: str | None = None_internal_call: bool = False )

将保存模型，以便您可以使用 from_pretrained() 重新加载它。

仅从主进程保存。

push_to_hub

< source >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

参数

commit_message (str，可选，默认为 "End of training") — 推送时提交的提交信息。
blocking (bool，可选，默认为 True) — 是否仅在 git push 完成后才返回。
token (str，可选，默认为 None) — 具有写入权限的 Token，用于覆盖 Trainer 的原始参数。
revision (str, 可选) — 要从中提交内容的 git 版本。默认为 “main” 分支的头部。
kwargs (dict[str, Any]，可选) — 传递给 ~Trainer.create_model_card 的额外关键字参数。

将 self.model 和 self.processing_class 上传到 🤗 模型库中的仓库 self.args.hub_model_id。

A2POConfig

类 trl.experimental.a2po.A2POConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 5e-05lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear'lr_scheduler_kwargs: dict | str | None = Nonewarmup_steps: float = 0optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused'optim_args: str | None = Noneweight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08optim_target_modules: None | str | list[str] = Nonegradient_accumulation_steps: int = 1average_tokens_across_devices: bool = Truemax_grad_norm: float = 1.0label_smoothing_factor: float = 0.0bf16: bool | None = Nonefp16: bool = Falsebf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: bool | None = Nonegradient_checkpointing: bool = Truegradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonetorch_compile: bool = Falsetorch_compile_backend: str | None = Nonetorch_compile_mode: str | None = Noneuse_liger_kernel: bool = Falseliger_kernel_config: dict[str, bool] | None = Noneuse_cache: bool = Falseneftune_noise_alpha: float | None = Nonetorch_empty_cache_steps: int | None = Noneauto_find_batch_size: bool = Falselogging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps'logging_steps: float = 10logging_first_step: bool = Falselog_on_each_node: bool = Truelogging_nan_inf_filter: bool = Trueinclude_num_input_tokens_seen: str | bool = 'no'log_level: str = 'passive'log_level_replica: str = 'warning'disable_tqdm: bool | None = Nonereport_to: None | str | list[str] = 'none'run_name: str | None = Noneproject: str = 'huggingface'trackio_space_id: str | None = Nonetrackio_bucket_id: str | None = Nonetrackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = Noneeval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no'eval_steps: float | None = Noneeval_delay: float = 0per_device_eval_batch_size: int = 8prediction_loss_only: bool = Falseeval_on_start: bool = Falseeval_do_concat_batches: bool = Trueeval_use_gather_object: bool = Falseeval_accumulation_steps: int | None = Noneinclude_for_metrics: list[str] = <factory>batch_eval_metrics: bool = Falsesave_only_model: bool = Falsesave_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps'save_steps: float = 500save_on_each_node: bool = Falsesave_total_limit: int | None = Noneenable_jit_checkpoint: bool = Falsepush_to_hub: bool = Falsehub_token: str | None = Nonehub_private_repo: bool | None = Nonehub_model_id: str | None = Nonehub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save'hub_always_push: bool = Falsehub_revision: str | None = Noneload_best_model_at_end: bool = Falsemetric_for_best_model: str | None = Nonegreater_is_better: bool | None = Noneignore_data_skip: bool = Falserestore_callback_states_from_checkpoint: bool = Falsefull_determinism: bool = Falseseed: int = 42data_seed: int | None = Noneuse_cpu: bool = Falseaccelerator_config: dict | str | None = Noneparallelism_config: accelerate.parallelism_config.ParallelismConfig | None = Nonedataloader_drop_last: bool = Falsedataloader_num_workers: int = 0dataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falsedataloader_prefetch_factor: int | None = Nonedataloader_multiprocessing_context: str | None = Nonedataloader_in_order: bool = Trueremove_unused_columns: bool = Falselabel_names: list[str] | None = Nonetrain_sampling_strategy: str = 'random'length_column_name: str = 'length'ddp_find_unused_parameters: bool | None = Noneddp_bucket_cap_mb: int | None = Noneddp_broadcast_buffers: bool | None = Noneddp_static_graph: bool | None = Noneddp_backend: str | None = Noneddp_timeout: int = 1800fsdp: str | None = Nonefsdp_config: dict[str, typing.Any] | str | None = Nonedeepspeed: dict | str | None = Nonedebug: str | list[transformers.debug_utils.DebugOption] = ''skip_memory_metrics: bool = Truedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseresume_from_checkpoint: str | None = Nonelocal_rank: int = -1model_init_kwargs: dict | None = Nonetrust_remote_code: bool = Falsemax_prompt_length: int | None = 512max_completion_length: int | None = 256temperature: float = 1.0top_p: float = 1.0top_k: int | None = Nonenum_value_samples: int = 8beta1: float = 0.5filter_all_incorrect: bool = Truebeta2: float = 0.001reward_weights: list[float] | None = None )

控制模型和参考模型的参数

model_init_kwargs (dict[str, Any], 可选) — 用于 from_pretrained 的关键字参数，当 experimental.a2po.A2POTrainer 的 model 参数以字符串形式提供时使用。revision 值在加载处理类时也会用到。
trust_remote_code (bool, optional, defaults to False) — Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to from_pretrained and from_pretrained.

控制数据预处理的参数

remove_unused_columns (bool，可选，默认为 False) — 是否仅保留数据集中的 "prompt" 列。如果您使用的自定义奖励函数需要 "prompts" 和 "completions" 以外的任何列，则应保持此参数为 False。

控制生成的参数

max_prompt_length (int 或 None，可选，默认为 512) — 提示词的最大长度。如果提示词超过此长度，将从左侧截断。
max_completion_length (int 或 None，可选，默认为 256) — 生成补全的最大长度。
temperature (float，可选，默认为 1.0) — 采样温度，用于第一阶段和第二阶段生成。
top_p (float，可选，默认为 1.0) — 用于控制所考虑的前几个 token 的累积概率的浮点数。取值必须在 (0, 1] 范围内。设置为 1.0 表示考虑所有 token。
top_k (int 或 None，可选) — 要保留的最高概率词汇表 token 数量。如果为 None，则禁用 top-k 过滤。

控制阶段 1（离线最优值估计）的参数

num_value_samples (int，可选，默认为 8) — 每个提示词从参考策略中采样的样本数量，用于估计 V*。
beta1 (float，可选，默认为 0.5) — 用于在第一阶段估计 V* 的 KL 温度。
filter_all_incorrect (bool，可选，默认为 True) — 是否丢弃所有参考样本均不正确的提示词。

控制 Stage 2（on-policy 回归）的参数

beta2 (float，可选，默认为 1e-3) — 在第二阶段回归目标中使用的 KL 温度。
reward_weights (list[float], 可选) — 每个奖励函数的权重。必须与奖励函数的数量相匹配。如果为 None，则所有奖励的权重均相等，权重为 1.0。

experimental.a2po.A2POTrainer 的配置类。

此类仅包含 A2PO 训练特有的参数。完整的训练参数列表，请参阅 TrainingArguments 文档。请注意，此类中的默认值可能与 TrainingArguments 中的不同。

Update on GitHub