Implementation:Hpcaitech ColossalAI Train KTO Script

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Preference Optimization, KTO, Distributed Training, RLHF
Last Updated	2026-02-09 00:00 GMT

Overview

train_kto.py is a training script that implements Kahneman-Tversky Optimization (KTO) for aligning language models using binary desirable/undesirable feedback signals.

Description

This script sets up a complete KTO training pipeline using ColossalAI's distributed training infrastructure. KTO is a preference optimization method that uses binary labels (desirable vs. undesirable) rather than paired preference data. The script initializes an actor model and a frozen reference model, configures distributed training with Gemini, ZeRO-2, or 3D hybrid parallelism plugins, and invokes the KTOTrainer to align the model. It validates that the desirable/undesirable weight ratio falls within the theoretical bounds from the KTO paper (Eq. 8 of arXiv:2402.01306), with optional automatic weight adjustment.

Usage

Use this script when you have binary preference data (desirable/undesirable labels) and want to align a language model using the KTO algorithm. It supports LoRA fine-tuning, gradient checkpointing, flash attention, and checkpoint resumption. Launch with torchrun for distributed execution.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalChat/examples/training_scripts/train_kto.py
• Lines: 1-379

Signature

def train(args) -> None

Import

# This is a standalone training script, typically run directly:
# torchrun --nproc_per_node=<N> train_kto.py --pretrain <model_path> --dataset <data_path>

I/O Contract

Inputs

Name	Type	Required	Description
--pretrain	str	Yes	Path to the pretrained model
--dataset	str (nargs=+)	Yes	Paths to tokenized training dataset(s)
--plugin	str	No	Plugin: gemini, gemini_auto, zero2, zero2_cpu, 3d (default: gemini)
--beta	float	No	Beta parameter in KTO loss (default: 0.1)
--desirable_weight	float	No	Weight for desirable samples in KTO loss (default: 1.0)
--undesirable_weight	float	No	Weight for undesirable samples in KTO loss (default: 1.0)
--auto_weight	flag	No	Automatically adjust weights to fit theoretical bounds
--eval_dataset	str (nargs=+)	No	Paths to evaluation dataset(s)
--checkpoint_path	str	No	Path to resume training from checkpoint
--lora_config	str	No	Path to LoRA configuration file
--max_length	int	No	Maximum sequence length (default: 2048)
--max_epochs	int	No	Maximum training epochs (default: 3)
--batch_size	int	No	Batch size per process (default: 4)
--lr	float	No	Learning rate (default: 5e-6)
--accumulation_steps	int	No	Gradient accumulation steps (default: 8)
--mixed_precision	str	No	Mixed precision: fp16 or bf16 (default: fp16)
--save_interval	int	No	Steps between checkpoints (default: 1000)
--grad_checkpoint	flag	No	Enable gradient checkpointing
--use_flash_attn	flag	No	Enable flash attention

Outputs

Name	Type	Description
checkpoint	directory	Final model checkpoint saved to --save_dir/modeling
config_file	JSON	Training configuration saved to --config_file

Usage Examples

# Train KTO with ZeRO-2 on 4 GPUs:
# torchrun --nproc_per_node=4 train_kto.py \
#     --pretrain meta-llama/Llama-2-7b \
#     --dataset ./kto_data \
#     --plugin zero2 \
#     --beta 0.1 \
#     --desirable_weight 1.0 \
#     --undesirable_weight 1.0 \
#     --auto_weight \
#     --lr 5e-6 \
#     --max_epochs 3 \
#     --save_dir ./kto_checkpoint

Key Features

• Weight Validation - Validates that the desirable/undesirable weight ratio satisfies the theoretical bounds [1, 4/3] from the KTO paper
• Automatic Weight Adjustment - Optional --auto_weight flag normalizes weights to fit within theoretical bounds
• Dual Booster Setup - Separate Booster instances for the actor model and frozen reference model
• LoRA Integration - Optional LoRA adaptation with weight merging at evaluation time
• Checkpoint Resumption - Supports loading from both full model checkpoints and training state checkpoints

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment