Implementation:OpenRLHF OpenRLHF Train KTO

Knowledge Sources	OpenRLHF KTO
Domains	Alignment, Training, CLI
Last Updated	2026-02-07 10:40 GMT

Overview

Concrete tool for launching KTO (Kahneman-Tversky Optimization) training from the command line.

Description

The train function is the entry point for KTO training. It orchestrates the full pipeline: initializes a DeepSpeed distributed strategy, loads the policy and reference models via the Actor class, prepares UnpairedPreferenceDataset with blended data sources, configures the optimizer and scheduler, and launches KTOTrainer.fit(). It supports LoRA/QLoRA fine-tuning, gradient checkpointing, checkpoint resumption, and both WandB and TensorBoard logging.

Usage

Use this entry point to launch KTO training from the command line. It is the primary user-facing interface for training with unpaired preference data and is invoked directly via python -m openrlhf.cli.train_kto or through DeepSpeed launcher.

Code Reference

Source Location

Repository: OpenRLHF
File: openrlhf/cli/train_kto.py
Lines: 1-272

Signature

def train(args) -> None:
    """
    KTO training entry point.

    Args:
        args: Namespace containing all training hyperparameters including
              pretrain, beta, desirable_loss_weight, undesirable_loss_weight,
              dataset, max_epochs, learning_rate, micro_train_batch_size, etc.
    """

Import

from openrlhf.cli.train_kto import train

I/O Contract

Inputs

Name	Type	Required	Description
args.pretrain	str	Yes	HuggingFace model name or path
args.dataset	str	Yes	Path to training dataset
args.beta	float	No	KTO regularization coefficient (default: 0.01)
args.desirable_loss_weight	float	No	Weight for desirable samples (default: 1.0)
args.undesirable_loss_weight	float	No	Weight for undesirable samples (default: 1.0)
args.learning_rate	float	No	Learning rate (default: 1e-5)
args.max_epochs	int	No	Training epochs (default: 1)
args.lora_rank	int	No	LoRA rank, 0 for full fine-tuning (default: 0)
args.save_path	str	No	Model save directory (default: ./ckpt)

Outputs

Name	Type	Description
Trained model	files	Saved to args.save_path in HuggingFace format
Checkpoints	files	Saved to args.ckpt_path during training
Logs	WandB/TensorBoard	Training metrics (kto_loss, chosen_reward, reject_reward)

Usage Examples

Launch KTO Training via DeepSpeed

deepspeed --module openrlhf.cli.train_kto \
    --save_path ./ckpt/llama3-8b-kto \
    --save_steps -1 \
    --logging_steps 1 \
    --eval_steps -1 \
    --micro_train_batch_size 8 \
    --train_batch_size 128 \
    --pretrain meta-llama/Meta-Llama-3-8B-Instruct \
    --learning_rate 1e-5 \
    --max_epochs 1 \
    --beta 0.01 \
    --dataset your/kto_dataset \
    --input_key input \
    --output_key output \
    --label_key label \
    --max_len 2048 \
    --zero_stage 2 \
    --apply_chat_template \
    --param_dtype bf16

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment