KTO 训练器

TRL 支持 Kahneman-Tversky 优化 (KTO) 训练器，用于使用二元反馈数据（例如，赞成/反对票）对齐语言模型，如 Kawin Ethayarajh、Winnie Xu、Niklas Muennighoff、Dan Jurafsky 和 Douwe Kiela 的论文中所述。有关完整示例，请查看 examples/scripts/kto.py。

根据您的基础模型的质量，您可能需要或不需要在 KTO 之前进行 SFT。这与标准的 RLHF 和 DPO 不同，后者始终需要 SFT。

预期数据集格式

KTO 训练器对数据集的格式有非常具体的要求，因为它不需要成对偏好。由于模型将经过训练以直接优化包含提示、模型补全和标签的示例，以指示补全结果是“好”还是“坏”，因此我们期望数据集具有以下列

提示
补全
标签

例如

kto_dataset_dict = {
    "prompt": [
        "Hey, hello",
        "How are you",
        "What is your name?",
        "What is your name?",
        "Which is the best programming language?",
        "Which is the best programming language?",
        "Which is the best programming language?",
    ],
    "completion": [
        "hi nice to meet you",
        "leave me alone",
        "I don't have a name",
        "My name is Mary",
        "Python",
        "C++",
        "Java",
    ],
    "label": [
        True,
        False,
        False,
        True,
        True,
        False,
        False,
    ],
}

其中 prompt 包含上下文输入，completion 包含相应的响应，label 包含相应的标志，指示生成的补全结果是期望的 (True) 还是不期望的 (False)。一个提示可以有多个响应，这反映在字典的值数组中重复的条目中。数据集必须至少包含一个期望的和至少一个不期望的补全结果。

预期模型格式

KTO 训练器期望使用 AutoModelForCausalLM 模型，而 PPO 则期望使用 AutoModelForCausalLMWithValueHead 模型用于价值函数。

使用 KTOTrainer

有关详细示例，请查看 examples/scripts/kto.py 脚本。在较高层面上，我们需要使用我们希望训练的 model 和参考 ref_model 初始化 KTOTrainer，我们将使用参考模型来计算首选和拒绝响应的隐式奖励。

beta 指的是隐式奖励的超参数，数据集包含上面列出的 3 个条目。请注意，model 和 ref_model 需要具有相同的架构（即仅解码器或编码器-解码器）。

desirable_weight 和 undesirable_weight 指的是对期望/正面示例和不期望/负面示例的损失施加的权重。默认情况下，它们都为 1。但是，如果您有更多的一种，那么您应该提高较不常见类型的权重，以使 (desirable_weight $\times$ 正面示例数量) 与 (undesirable_weight $\times$ 负面示例数量) 的比率在 1:1 到 4:3 的范围内。

强烈建议您使用介于 `5e-7` 和 `5e-6` 之间的学习率，有效批次大小介于 `8` 和 `32` 之间，用于 LoRA 和完全微调。即使您正在处理小型数据集，我们也不建议使用此范围之外的学习率；相反，使用较小的批次大小和/或更多的训练 epoch 将为您带来更好的结果。

training_args = KTOConfig(
    beta=0.1,
    desirable_weight=1.0,
    undesirable_weight=1.0,
    learning_rate=5e-7,
)

kto_trainer = KTOTrainer(
    model,
    ref_model,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)

在此之后，可以调用

kto_trainer.train()

对于专家混合模型：启用辅助损失

如果专家之间的负载大致均匀分布，则 MOE 最有效。
为了确保我们在偏好调整期间以类似的方式训练 MOE，最好将负载均衡器的辅助损失添加到最终损失中。

通过在模型配置（例如 MixtralConfig）中设置 output_router_logits=True 来启用此选项。
要缩放辅助损失对总损失的贡献程度，请使用超参数 router_aux_loss_coef=...（默认值：0.001）。

KTOTrainer

class trl.KTOTrainer

< 源代码 >

( model: Union = None ref_model: Union = None args: KTOConfig = None train_dataset: Optional = None eval_dataset: Union = None tokenizer: Optional = None data_collator: Optional = None model_init: Optional = None callbacks: Optional = None optimizers: Tuple = (None, None) preprocess_logits_for_metrics: Optional = None peft_config: Optional = None compute_metrics: Optional = None model_adapter_name: Optional = None ref_adapter_name: Optional = None )

参数

model (transformers.PreTrainedModel) — 要训练的模型，最好是 AutoModelForSequenceClassification。
ref_model (PreTrainedModelWrapper) — 具有因果语言建模头的 Hugging Face transformer 模型。用于隐式奖励计算和损失。如果未提供参考模型，则训练器将创建一个与要优化的模型具有相同架构的参考模型。
args (KTOConfig) — 用于训练的参数。
train_dataset (datasets.Dataset) — 用于训练的数据集。
eval_dataset (datasets.Dataset) — 用于评估的数据集。
tokenizer (transformers.PreTrainedTokenizerBase) — 用于训练的分词器。如果您想使用默认数据整理器，则此参数是必需的。
data_collator (transformers.DataCollator, 可选, 默认为 None) — 用于训练的数据整理器。如果未指定 None，则将使用默认数据整理器 (DPODataCollatorWithPadding)，它将成对序列的数据集填充到批次中最长序列的长度。
model_init (Callable[[], transformers.PreTrainedModel]) — 用于训练的模型初始化器。如果未指定 None，则将使用默认模型初始化器。
callbacks (List[transformers.TrainerCallback]) — 用于训练的回调。
optimizers (Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]) — 用于训练的优化器和调度器。
preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]) — 用于预处理 logits 以便计算指标的函数。
peft_config (Dict, defaults to None) — 用于训练的 PEFT 配置。如果传递 PEFT 配置，模型将被封装在 PEFT 模型中。
disable_dropout (bool, defaults to True) — 是否禁用 model 和 ref_model 中的 dropout。
compute_metrics (Callable[[EvalPrediction], Dict], optional) — 用于计算指标的函数。必须接受一个 EvalPrediction 并返回一个从字符串到指标值的字典。
model_adapter_name (str, defaults to None) — 当使用带有多个适配器的 LoRA 时，训练目标 PEFT 适配器的名称。
ref_adapter_name (str, defaults to None) — 当使用带有多个适配器的 LoRA 时，参考 PEFT 适配器的名称。

初始化 KTOTrainer。

compute_reference_log_probs

< source >

( padded_batch: Dict )

计算参考模型对于 KTO 特定数据集的单个填充批次的对数概率。

evaluation_loop

< source >

( dataloader: DataLoader description: str prediction_loss_only: Optional = None ignore_keys: Optional = None metric_key_prefix: str = 'eval' )

重写内置的评估循环，以存储每个批次的指标。 Trainer.evaluate() 和 Trainer.predict() 共享的预测/评估循环。

有标签或无标签均可使用。

get_batch_logps

< source >

( logits: FloatTensor labels: LongTensor average_log_prob: bool = False label_pad_token_id: int = -100 is_encoder_decoder: bool = False )

计算给定 logits 下给定标签的对数概率。

get_batch_loss_metrics

< source >

( model batch: Dict )

计算给定输入批次的 KTO 损失和其他指标，用于训练或测试。

get_batch_samples

< source >

( model batch: Dict )

为给定的输入批次从模型和参考模型生成样本。

get_eval_dataloader

< source >

( eval_dataset: Optional = None )

参数

eval_dataset (torch.utils.data.Dataset, optional) — 如果提供，将覆盖 self.eval_dataset。如果它是一个 Dataset，则会自动删除 model.forward() 方法不接受的列。它必须实现 __len__。

返回评估 ~torch.utils.data.DataLoader。

transformers.src.transformers.trainer.get_eval_dataloader 的子类，用于预计算 ref_log_probs。

get_train_dataloader

< source >

( )

返回训练 ~torch.utils.data.DataLoader。

transformers.src.transformers.trainer.get_train_dataloader 的子类，用于预计算 ref_log_probs。

kto_loss

< source >

( policy_chosen_logps: FloatTensor policy_rejected_logps: FloatTensor policy_KL_logps: FloatTensor reference_chosen_logps: FloatTensor reference_rejected_logps: FloatTensor reference_KL_logps: FloatTensor ) → 四个张量的元组

四个张量的元组

(losses, chosen_rewards, rejected_rewards, KL)。 losses 张量包含批次中每个示例的 KTO 损失。 chosen_rewards 和 rejected_rewards 张量分别包含所选和拒绝响应的奖励。 KL 张量包含策略模型和参考模型之间分离的 KL 散度估计。

计算一批策略和参考模型对数概率的 KTO 损失。

log

< source >

( logs: Dict )

参数

logs (Dict[str, float]) — 要记录的值。

在各种监视训练的对象（包括存储的指标）上记录 logs。

null_ref_context

< source >

( )

用于处理空参考模型（即 PEFT 适配器操作）的上下文管理器。

KTOConfig

class trl.KTOConfig

< source >

( output_dir: str overwrite_output_dir: bool = False do_train: bool = False do_eval: bool = False do_predict: bool = False eval_strategy: Union = '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: Optional = None per_gpu_eval_batch_size: Optional = None gradient_accumulation_steps: int = 1 eval_accumulation_steps: Optional = None eval_delay: Optional = 0 torch_empty_cache_steps: Optional = None learning_rate: float = 5e-07 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: Union = 'linear' lr_scheduler_kwargs: Union = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: Optional = 'passive' log_level_replica: Optional = 'warning' log_on_each_node: bool = True logging_dir: Optional = None logging_strategy: Union = 'steps' logging_first_step: bool = False logging_steps: float = 500 logging_nan_inf_filter: bool = True save_strategy: Union = 'steps' save_steps: float = 500 save_total_limit: Optional = None save_safetensors: Optional = 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: Optional = 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: Optional = None local_rank: int = -1 ddp_backend: Optional = None tpu_num_cores: Optional = None tpu_metrics_debug: bool = False debug: Union = '' dataloader_drop_last: bool = False eval_steps: Optional = None dataloader_num_workers: int = 0 dataloader_prefetch_factor: Optional = None past_index: int = -1 run_name: Optional = None disable_tqdm: Optional = None remove_unused_columns: Optional = True label_names: Optional = None load_best_model_at_end: Optional = False metric_for_best_model: Optional = None greater_is_better: Optional = None ignore_data_skip: bool = False fsdp: Union = '' fsdp_min_num_params: int = 0 fsdp_config: Union = None fsdp_transformer_layer_cls_to_wrap: Optional = None accelerator_config: Union = None deepspeed: Union = None label_smoothing_factor: float = 0.0 optim: Union = 'adamw_torch' optim_args: Optional = None adafactor: bool = False group_by_length: bool = False length_column_name: Optional = 'length' report_to: Union = None ddp_find_unused_parameters: Optional = None ddp_bucket_cap_mb: Optional = None ddp_broadcast_buffers: Optional = 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: Optional = None hub_model_id: Optional = None hub_strategy: Union = 'every_save' hub_token: Optional = None hub_private_repo: bool = False hub_always_push: bool = False gradient_checkpointing: bool = False gradient_checkpointing_kwargs: Union = None include_inputs_for_metrics: bool = False include_for_metrics: List = <factory> eval_do_concat_batches: bool = True fp16_backend: str = 'auto' evaluation_strategy: Union = None push_to_hub_model_id: Optional = None push_to_hub_organization: Optional = None push_to_hub_token: Optional = None mp_parameters: str = '' auto_find_batch_size: bool = False full_determinism: bool = False torchdynamo: Optional = None ray_scope: Optional = 'last' ddp_timeout: Optional = 1800 torch_compile: bool = False torch_compile_backend: Optional = None torch_compile_mode: Optional = None dispatch_batches: Optional = None split_batches: Optional = None include_tokens_per_second: Optional = False include_num_input_tokens_seen: Optional = False neftune_noise_alpha: Optional = None optim_target_modules: Union = None batch_eval_metrics: bool = False eval_on_start: bool = False use_liger_kernel: Optional = False eval_use_gather_object: Optional = False max_length: Optional = None max_prompt_length: Optional = None max_completion_length: Optional = None beta: float = 0.1 loss_type: Literal = 'kto' desirable_weight: float = 1.0 undesirable_weight: float = 1.0 label_pad_token_id: int = -100 padding_value: Optional = None truncation_mode: str = 'keep_end' generate_during_eval: bool = False is_encoder_decoder: Optional = None precompute_ref_log_probs: bool = False model_init_kwargs: Optional = None ref_model_init_kwargs: Optional = None dataset_num_proc: Optional = None )

参数

learning_rate (float, optional, defaults to 5e-7) — Initial learning rate for AdamW optimizer. The default value replaces that of TrainingArguments.
max_length (Optional[int], optional, defaults to None) — Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want to use the default data collator.
max_prompt_length (Optional[int], optional, defaults to None) — Maximum length of the prompt. This argument is required if you want to use the default data collator.
max_completion_length (Optional[int], optional, defaults to None) — Maximum length of the completion. This argument is required if you want to use the default data collator and your model is an encoder-decoder.
beta (float, optional, defaults to 0.1) — Parameter controlling the deviation from the reference model. Higher β means less deviation from the reference model.
loss_type (str, optional, defaults to "kto") — Type of loss to use. Possible values are:
- "kto": KTO loss from the KTO paper.
- "apo_zero_unpaired": Unpaired variant of APO-zero loss from the APO paper.
undesirable_weight (float, optional, defaults to 1.0) — Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs.
label_pad_token_id (int, 可选, 默认为 -100) — 标签填充 token ID。如果你想使用默认数据收集器，则此参数是必需的。
padding_value (Optional[int], 可选, 默认为 None) — 要使用的填充值。如果为 None，则使用 tokenizer 的填充值。
truncation_mode (str, 可选, 默认为 "keep_end") — 当 prompt 过长时使用的截断模式。可能的值为 "keep_end" 或 "keep_start"。如果你想使用默认数据收集器，则此参数是必需的。
generate_during_eval (bool, 可选, 默认为 False) — 如果为 True，则在评估期间从模型和参考模型生成补全内容并记录到 W&B。
is_encoder_decoder (Optional[bool], 可选, 默认为 None) — 当使用 model_init 参数（可调用对象）来实例化模型而不是 model 参数时，你需要指定可调用对象返回的模型是否为 encoder-decoder 模型。
precompute_ref_log_probs (bool, 可选, 默认为 False) — 是否为训练和评估数据集预计算参考模型对数概率。当在没有参考模型的情况下进行训练以减少所需的总 GPU 内存时，这非常有用。
model_init_kwargs (Optional[Dict[str, Any]], 可选, 默认为 None) — 当从字符串实例化模型时，传递给 AutoModelForCausalLM.from_pretrained 的关键字参数。
ref_model_init_kwargs (Optional[Dict[str, Any]], 可选, 默认为 None) — 当从字符串实例化参考模型时，传递给 AutoModelForCausalLM.from_pretrained 的关键字参数。 dataset_num_proc — (Optional[int], 可选, 默认为 None): 用于处理数据集的进程数。

用于 KTOTrainer 的配置类。

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

< > 在 GitHub 上更新