explain what this large language trainer is and basics steps to implement it alongwith simple examples (with easy to understand explanations): ""Hugging Face's logo
Hugging Face
Search models, datasets, users...

TRL documentation

KTO Trainer


You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.8.1).
KTO Trainer
TRL supports the Kahneman-Tversky Optimization (KTO) Trainer for aligning language models with binary feedback data (e.g., upvote/downvote), as described in the paper by Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. For a full example have a look at examples/scripts/kto.py.

Depending on how good your base model is, you may or may not need to do SFT before KTO. This is different from standard RLHF and DPO, which always require SFT.

Expected dataset format
The KTO trainer expects a very specific format for the dataset as it does not require pairwise preferences. Since the model will be trained to directly optimize examples that consist of a prompt, model completion, and a label to indicate whether the completion is “good” or “bad”, we expect a dataset with the following columns:

prompt
completion
label
for example:

Copied
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,
    ],
}
where the prompt contains the context inputs, completion contains the corresponding responses and label contains the corresponding flag that indicates if the generated completion is desired (True) or undesired (False). A prompt can have multiple responses and this is reflected in the entries being repeated in the dictionary’s value arrays.

Expected model format
The KTO trainer expects a model of AutoModelForCausalLM, compared to PPO that expects AutoModelForCausalLMWithValueHead for the value function.

Using the KTOTrainer
For a detailed example have a look at the examples/scripts/kto.py script. At a high level we need to initialize the KTOTrainer with a model we wish to train and a reference ref_model which we will use to calculate the implicit rewards of the preferred and rejected response.

The beta refers to the hyperparameter of the implicit reward, and the dataset contains the 3 entries listed above. Note that the model and ref_model need to have the same architecture (ie decoder only or encoder-decoder).

The desirable_weight and undesirable_weight refer to the weights placed on the losses for desirable/positive and undesirable/negative examples. By default, they are both 1. However, if you have more of one or the other, then you should upweight the less common type such that the ratio of (desirable_weight number of positives) to (undesirable_weight number of negatives) is in the range 1:1 to 4:3.

Copied
training_args = KTOConfig(
    beta=0.1,
    desirable_weight=1.0,
    undesirable_weight=1.0,
)

kto_trainer = KTOTrainer(
    model,
    model_ref,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)
After this one can then call:

Copied
kto_trainer.train()
KTOTrainer
class trl.KTOTrainer
<
source
>
( model: Union = Noneref_model: Union = Noneargs: KTOConfig = Nonetrain_dataset: Optional = Noneeval_dataset: Union = Nonetokenizer: Optional = Nonedata_collator: Optional = Nonemodel_init: Optional = Nonecallbacks: Optional = Noneoptimizers: Tuple = (None, None)preprocess_logits_for_metrics: Optional = Nonepeft_config: Optional = Nonecompute_metrics: Optional = None )

Expand 14 parameters
Parameters

model (transformers.PreTrainedModel) — The model to train, preferably an AutoModelForSequenceClassification.
ref_model (PreTrainedModelWrapper) — Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation and loss. If no reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized.
args (KTOConfig) — The arguments to use for training.
train_dataset (datasets.Dataset) — The dataset to use for training.
eval_dataset (datasets.Dataset) — The dataset to use for evaluation.
tokenizer (transformers.PreTrainedTokenizerBase) — The tokenizer to use for training. This argument is required if you want to use the default data collator.
data_collator (transformers.DataCollator, optional, defaults to None) — The data collator to use for training. If None is specified, the default data collator (DPODataCollatorWithPadding) will be used which will pad the sequences to the maximum length of the sequences in the batch, given a dataset of paired sequences.
model_init (Callable[[], transformers.PreTrainedModel]) — The model initializer to use for training. If None is specified, the default model initializer will be used.
callbacks (List[transformers.TrainerCallback]) — The callbacks to use for training.
optimizers (Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]) — The optimizer and scheduler to use for training.
preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]) — The function to use to preprocess the logits before computing the metrics.
peft_config (Dict, defaults to None) — The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model.
disable_dropout (bool, defaults to True) — Whether or not to disable dropouts in model and ref_model.
compute_metrics (Callable[[EvalPrediction], Dict], optional) — The function to use to compute the metrics. Must take a EvalPrediction and return a dictionary string to metric values.
Initialize KTOTrainer.

build_tokenized_answer
<
source
>
( promptanswer )

Llama tokenizer does not satisfy enc(a + b) = enc(a) + enc(b). It does ensure enc(a + b) = enc(a) + enc(a + b)[len(enc(a)):]. Reference: https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257

compute_reference_log_probs
<
source
>
( padded_batch: Dict )

Computes log probabilities of the reference model for a single padded batch of a KTO specific dataset.

evaluation_loop
<
source
>
( dataloader: DataLoaderdescription: strprediction_loss_only: Optional = Noneignore_keys: Optional = Nonemetric_key_prefix: str = 'eval' )

Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().

Works both with or without labels.

get_KL_dataset
<
source
>
( batch )

Creates mismatched pairs of prompts and completions for the KL dataset.

get_batch_logps
<
source
>
( logits: FloatTensorlabels: LongTensoraverage_log_prob: bool = Falselabel_pad_token_id: int = -100is_encoder_decoder: bool = False )

Compute the log probabilities of the given labels under the given logits.

get_batch_loss_metrics
<
source
>
( modelbatch: Dicttrain_eval: Literal = 'train' )

Compute the KTO loss and other metrics for the given batch of inputs for train or test.

get_batch_samples
<
source
>
( modelbatch: Dict )

Generate samples from the model and reference model for the given batch of inputs.

get_eval_dataloader
<
source
>
( eval_dataset: Optional = None )

Parameters

eval_dataset (torch.utils.data.Dataset, optional) — If provided, will override self.eval_dataset. If it is a Dataset, columns not accepted by the model.forward() method are automatically removed. It must implement __len__.
Returns the evaluation ~torch.utils.data.DataLoader.

Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute ref_log_probs.

get_train_dataloader
<
source
>
( )

Returns the training ~torch.utils.data.DataLoader.

Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute ref_log_probs.

kto_loss
<
source
>
( policy_chosen_logps: FloatTensorpolicy_rejected_logps: FloatTensorpolicy_KL_logps: FloatTensorreference_chosen_logps: FloatTensorreference_rejected_logps: FloatTensorreference_KL_logps: FloatTensor ) → A tuple of four tensors

Returns

A tuple of four tensors

(losses, chosen_rewards, rejected_rewards, KL). The losses tensor contains the KTO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively. The KL tensor contains the detached KL divergence estimate between the policy and reference models.

Compute the KTO loss for a batch of policy and reference model log probabilities.

log
<
source
>
( logs: Dict )

Parameters

logs (Dict[str, float]) — The values to log.
Log logs on the various objects watching training, including stored metrics.

tokenize_row
<
source
>
( featuremodel: Union = Noneprefix = '' )

Tokenize a single row from a KTO specific dataset.

At this stage, we don’t convert to PyTorch tensors yet; we just handle the truncation in case the prompt + completion responses is/are too long. First we truncate the prompt; if we’re still too long, we truncate the completion.

We also create the labels for the completion responses, which are of length equal to the sum of the length of the prompt and the completion response, with label_pad_token_id for the prompt tokens.

KTOConfig
class trl.KTOConfig
<
source
>
( output_dir: stroverwrite_output_dir: bool = Falsedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseevaluation_strategy: Union = 'no'prediction_loss_only: bool = Falseper_device_train_batch_size: int = 8per_device_eval_batch_size: int = 8per_gpu_train_batch_size: Optional = Noneper_gpu_eval_batch_size: Optional = Nonegradient_accumulation_steps: int = 1eval_accumulation_steps: Optional = Noneeval_delay: Optional = 0learning_rate: float = 5e-05weight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08max_grad_norm: float = 1.0num_train_epochs: float = 3.0max_steps: int = -1lr_scheduler_type: Union = 'linear'lr_scheduler_kwargs: Optional = <factory>warmup_ratio: float = 0.0warmup_steps: int = 0log_level: Optional = 'passive'log_level_replica: Optional = 'warning'log_on_each_node: bool = Truelogging_dir: Optional = Nonelogging_strategy: Union = 'steps'logging_first_step: bool = Falselogging_steps: float = 500logging_nan_inf_filter: bool = Truesave_strategy: Union = 'steps'save_steps: float = 500save_total_limit: Optional = Nonesave_safetensors: Optional = Truesave_on_each_node: bool = Falsesave_only_model: bool = Falseno_cuda: bool = Falseuse_cpu: bool = Falseuse_mps_device: bool = Falseseed: int = 42data_seed: Optional = Nonejit_mode_eval: bool = Falseuse_ipex: bool = Falsebf16: bool = Falsefp16: bool = Falsefp16_opt_level: str = 'O1'half_precision_backend: str = 'auto'bf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: Optional = Nonelocal_rank: int = -1ddp_backend: Optional = Nonetpu_num_cores: Optional = Nonetpu_metrics_debug: bool = Falsedebug: Union = ''dataloader_drop_last: bool = Falseeval_steps: Optional = Nonedataloader_num_workers: int = 0dataloader_prefetch_factor: Optional = Nonepast_index: int = -1run_name: Optional = Nonedisable_tqdm: Optional = Noneremove_unused_columns: Optional = Truelabel_names: Optional = Noneload_best_model_at_end: Optional = Falsemetric_for_best_model: Optional = Nonegreater_is_better: Optional = Noneignore_data_skip: bool = Falsefsdp: Union = ''fsdp_min_num_params: int = 0fsdp_config: Union = Nonefsdp_transformer_layer_cls_to_wrap: Optional = Noneaccelerator_config: Optional = Nonedeepspeed: Optional = Nonelabel_smoothing_factor: float = 0.0optim: Union = 'adamw_torch'optim_args: Optional = Noneadafactor: bool = Falsegroup_by_length: bool = Falselength_column_name: Optional = 'length'report_to: Optional = Noneddp_find_unused_parameters: Optional = Noneddp_bucket_cap_mb: Optional = Noneddp_broadcast_buffers: Optional = Nonedataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falseskip_memory_metrics: bool = Trueuse_legacy_prediction_loop: bool = Falsepush_to_hub: bool = Falseresume_from_checkpoint: Optional = Nonehub_model_id: Optional = Nonehub_strategy: Union = 'every_save'hub_token: Optional = Nonehub_private_repo: bool = Falsehub_always_push: bool = Falsegradient_checkpointing: bool = Falsegradient_checkpointing_kwargs: Optional = Noneinclude_inputs_for_metrics: bool = Falsefp16_backend: str = 'auto'push_to_hub_model_id: Optional = Nonepush_to_hub_organization: Optional = Nonepush_to_hub_token: Optional = Nonemp_parameters: str = ''auto_find_batch_size: bool = Falsefull_determinism: bool = Falsetorchdynamo: Optional = Noneray_scope: Optional = 'last'ddp_timeout: Optional = 1800torch_compile: bool = Falsetorch_compile_backend: Optional = Nonetorch_compile_mode: Optional = Nonedispatch_batches: Optional = Nonesplit_batches: Optional = Noneinclude_tokens_per_second: Optional = Falseinclude_num_input_tokens_seen: Optional = Falseneftune_noise_alpha: Optional = Noneoptim_target_modules: Union = Nonemax_length: Optional = Nonemax_prompt_length: Optional = Nonemax_completion_length: Optional = Nonebeta: float = 0.1desirable_weight: Optional = 1.0undesirable_weight: Optional = 1.0label_pad_token_id: int = -100padding_value: int = Nonetruncation_mode: str = 'keep_end'generate_during_eval: bool = Falseis_encoder_decoder: Optional = Noneprecompute_ref_log_probs: bool = Falsemodel_init_kwargs: Optional = Noneref_model_init_kwargs: Optional = None )

Expand 12 parameters
Parameters

max_length (int, optional, defaults to None) — The maximum length of the sequences in the batch. This argument is required if you want to use the default data collator.
max_prompt_length (int, optional, defaults to None) — The maximum length of the prompt. This argument is required if you want to use the default data collator.
max_completion_length (int, optional, defaults to None) — The maximum length of the target. This argument is required if you want to use the default data collator and your model is an encoder-decoder.
beta (float, defaults to 0.1) — The beta factor in KTO loss. Higher beta means less divergence from the initial policy.
desirable_weight (float, optional, defaults to 1.0) — The desirable losses are weighed by this factor to counter unequal number of desirable and undesirable paris.
undesirable_weight (float, optional, defaults to 1.0) — The undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs.
label_pad_token_id (int, defaults to -100) — The label pad token id. This argument is required if you want to use the default data collator.
padding_value (int, defaults to 0) — The padding value if it is different to the tokenizer’s pad_token_id.
truncation_mode (str, defaults to keep_end) — The truncation mode to use, either keep_end or keep_start. This argument is required if you want to use the default data collator.
generate_during_eval (bool, defaults to False) — Whether to sample and log generations during evaluation step.
is_encoder_decoder (Optional[bool], optional, defaults to None) — If no model is provided, we need to know if the model_init returns an encoder-decoder.
precompute_ref_log_probs (bool, defaults to False) — Flag to precompute reference model log probabilities for training and evaluation datasets. This is useful if you want to train without the reference model and reduce the total GPU memory needed. model_init_kwargs — (Optional[Dict], optional): Dict of Optional kwargs to pass when instantiating the model from a string. ref_model_init_kwargs — (Optional[Dict], optional): Dict of Optional kwargs to pass when instantiating the ref model from a string.
KTOConfig collects all training arguments related to the KTOTrainer class.

Using HfArgumentParser we can turn this class into argparse arguments that can be specified on the command line.

←
DPO Trainer
CPO Trainer
→""