Add KTO Loss #475
Add KTO Loss #475
Conversation
There was a problem hiding this comment.
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 About |
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> ### KTO LOSS #### Memory  #### Speed  <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence
|
AMD Test failed due to no gpu available, not related to the PR: |
| H: int, | ||
| V: int, | ||
| dtype: torch.dtype, | ||
| bias: bool = False, |
There was a problem hiding this comment.
Seems like bias is being used as a boolean, whereas self.KTO_loss requires it to be an optional tensor. Can we modify the param names to avoid confusion?
| preference_labels_chunk=None, | ||
| ref_input_chunk=None, | ||
| ): | ||
| (chunk_grad_input, chunk_grad_weight, *chunk_grad_bias), (chunk_loss) = fused_fwd_bwd( |
There was a problem hiding this comment.
Why is it *chunk_grad_bias and not chunk_grad_bias like the other gradients?
| input_chunk = input_chunk | ||
| ref_input_chunk = ref_input_chunk if use_ref_model else None | ||
| target_chunk = target_chunk |
There was a problem hiding this comment.
the reassignment can be avoided.
| ref_input_chunk = ref_input_chunk if use_ref_model else None | ||
| target_chunk = target_chunk | ||
|
|
||
| # mark input_chunk, target_chunk, and target dimension 1 as dynamic to prevent torch.compile recompilation |
There was a problem hiding this comment.
let's mention that the dimension 1 corresponds to sequence length.
| """ | ||
| Compute the total loss for a chunk of input and target, while using an alignment/preference loss function. | ||
| Args: | ||
| preference_loss_fn (callable): Loss function to compute the loss on a chunk of input/target. |
There was a problem hiding this comment.
Seems like this class already has a staticmethod for preference_loss_fn. Why do we need an extra arg here?
Summary
Close KTO Item of the Roadmap: #371
Implements the Kahneman-Tversky Optimization (KTO) loss function.
KTO Loss Function
For a policy π compared to a reference policy π₀:
When y is chosen:
When y is rejected:
where:
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:
Key Changes
LigerFusedLinearKTOLossclassLigerFusedLinearKTOFunctionfor the core KTO computationtest_kto_loss.pyHFKTOLoss) based on Hugging Face's implementationReference
Testing Done
Test is passing now:
pytest test/chunked_loss/test_kto_loss.pymake testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence