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

Special thanks to @shivam15s on the optimization PR: #491, otherwise my implementation won't achieve speed as list below

Memory:

Speed:

Notable learning on optimizing the speed:

• [Culprit] Repeated calculation of KL when we split to N chunks
• [Good to have] Remove the unnecessary variables calculation like aux_outputs

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()

[hebiao064]

AMD Test failed due to no gpu available, not related to the PR: FAILED test/transformers/test_swiglu.py::test_correctness_functional[dtype1-10000.0-0.01-9-7-41] - RuntimeError: No HIP GPUs are available

[kvignesh1420]

Thanks for the contribution. Did a first pass over the functionality and left some comments.

[kvignesh1420]

Why is it *chunk_grad_bias and not chunk_grad_bias like the other gradients?

[hebiao064]

The *chunk_grad_bias syntax is used here because of how Python handles unpacking in this situation. Let me explain:

1. The fused_fwd_bwd function returns a tuple of two elements:

   • First element: A tuple of gradients (either 2 or 3 elements depending on if bias is used)
   • Second element: The loss value

2. When bias is None, fused_fwd_bwd returns only two gradients: (chunk_grad_input, chunk_grad_weight)

   When bias is not None, it returns three: (chunk_grad_input, chunk_grad_weight, chunk_grad_bias)

3. The * operator is used to handle this variability - it collects any remaining elements into a list. So:

   • If bias is None: *chunk_grad_bias becomes an empty list []
   • If bias exists: *chunk_grad_bias becomes a list with one element [chunk_grad_bias]

[kvignesh1420]

Seems like this class already has a staticmethod for preference_loss_fn. Why do we need an extra arg here?

[hebiao064]

1. The abstract preference_loss_fn method defined at the class level is just a placeholder that defines the interface - notice it's marked with @abstractmethod and raises NotImplementedError. This is meant to be implemented by subclasses.

2. The preference_loss_fn parameter in the methods like forward() and _compute_loss() is the actual function that will be used to compute the loss. This is passed in at runtime.

This is a common pattern in Python where the base class defines an interface (abstract method) but allows for runtime flexibility by accepting the actual implementation as a parameter. This makes the class more flexible as you can:

1. Create subclasses that implement a default preference loss function by overriding the abstract method
2. Override that default at runtime by passing in a different loss function as a parameter