Add KTO Loss

[hebiao064]

Summary

Close KTO Item of the Roadmap: #371

Implements the Kahneman-Tversky Optimization (KTO) loss function.

• Paper: "KTO: Model Alignment as Prospect Theory-Guided Optimization"

KTO Loss Function

For a policy π compared to a reference policy π₀:

When y is chosen:

$L_{KTO} = 1 - \sigma(\beta \cdot (\log[\frac{\pi(x)}{\pi_0(x)}] - KL(\pi||\pi_0)_y))$

When y is rejected:

$L_{KTO} = 1 - \sigma(\beta \cdot (KL(\pi||\pi_0)_y - \log[\frac{\pi(x)}{\pi_0(x)}]))$

where:

• σ is the sigmoid function
• β is a temperature parameter
• KL(π||π₀)_y is the KL divergence threshold for action y

Intuition

KTO loss is inspired by prospect theory from behavioral economics, which models how humans make decisions under uncertainty.

The loss function is asymmetric, treating gains and losses differently, similar to
human decision-making patterns.

Credit by: https://www.youtube.com/watch?v=nSrj1J6ODoM&t=422s

Benchmark Result

Memory:

Speed:

Key Changes

• Implemented LigerFusedLinearKTOLoss class
• Added LigerFusedLinearKTOFunction for the core KTO computation
• Created comprehensive test suite in test_kto_loss.py
• Added reference implementation (HFKTOLoss) based on Hugging Face's implementation

Reference

• Based on Hugging Face's implementation: https://github.com/huggingface/trl/blob/main/trl/trainer/kto_trainer.py

Testing Done

Test is passing now:
pytest test/chunked_loss/test_kto_loss.py

• Parameterized tests covering various configurations:
  • Different batch sizes, sequence lengths, hidden dims, and vocab sizes
  • Multiple data types (bfloat16, float32)
  • Bias and reference bias variations
  • Different ignore indices and beta values
• Correctness tests comparing against reference implementation
• Gradient checking and backward pass verification

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[ByronHsu]

Take a brief look, I am not very familiar with KTO math but why do we not have KL_log_probs but original HF has https://github.com/huggingface/trl/blob/cd7156fb34ddf9a8c04fcd640a4067933461d44e/trl/trainer/kto_trainer.py#L1121. We also need to be careful about scaling. Seems in original HF, kto_loss returns an unreduced version, but we probably need to reduce as mean. cc @shivam15s

[hebiao064]

Take a brief look, I am not very familiar with KTO math but why do we not have KL_log_probs but original HF has https://github.com/huggingface/trl/blob/cd7156fb34ddf9a8c04fcd640a4067933461d44e/trl/trainer/kto_trainer.py#L1121. We also need to be careful about scaling. Seems in original HF, kto_loss returns an unreduced version, but we probably need to reduce as mean. cc @shivam15s

About KL, I'll take a further look in trl about how to support that.

About reduce, HF did averaged it here: loss = losses.nanmean()