Adding a KTO Optimizer

[bjb19]

Due to recent changes in the space we probably need to take a look at adding this similarly to our DPO recipes, Happy to take a stab at it we are in agreement.

[krammnic]

Probably then ORPO and other DPO-like losses should be added?

[krammnic]

Actually, if some preference optimization methods will be added, this dpo file should be really refactored. Maybe there should be something like:

dpo/method1.py
dpo/method2.py
...

[SalmanMohammadi]

Hey @bjb19 @krammnic! Always happy to see interest in our RLHF offerings.

Due to recent changes in the space

As someone who is embarrasingly out-of-the-loop, I'm curious to hear what kind of recent changes you had in mind?

Happy to take a stab at it we are in agreement

From the last time I looked at it I think it should be straightforward enough to add as a new loss to our DPO recipe, but I can double check. It'd be really helpful as a first step to scope this out a bit more, e.g. any motivation for the technique/how a reference implementation works/how this would work in torchtune.

EDIT: OK at first glance this may require a bit more thought as TRL's KTOTrainer is quite different to its DPOTrainer.

We try to ensure all of the methods we support in torchtune are as well-tested as they can be. When implementing a new finetuning technique, we usually like to see parity with an existing implementation, e.g. by comparing identical training runs using a reference implementation such as TRL. There's an example PR here for SimPO, and I think the process for contributing KTO would be pretty similar.

If you're interested in taking this on I'd be more than happy to support you.

ORPO and other DPO-like losses should be added

I'm here for more DPO losses! We actually have an outstanding issue for IPO (#1291) and CPO (#1290). If there's specific techniques you'd like to see a great place to start would be opening an issue to keep track of them and kick off discussion on how to proceed.

Actually, if some preference optimization methods will be added, this dpo file should be really refactored. Maybe there should be something like:

This is a pretty good point! It's something we deliberated a bit over. One of the core tenants in the torchtune design philosophy is simplicity and extensibility - favouring code duplication over abstractions. I think there may come a time when we look at adding a new technique to the DPO recipe and we start to think, "OK, having X different branches for Y different techniques is just getting ugly, we need to refactor this". Currently, the different techniques we support in DPO differ in such a minor way that it doesn't make a huge amount of sense to write an entirely new recipe file to modify a couple lines. I think this is definitely something to keep in mind as we look towards extending our DPO-esque offerings to ensure the recipes remain as clean and usable as they can be.

[krammnic]

@SalmanMohammadi Can we then determine which losses would be interesting to add and create some general issue about them with some progress traking?

[SalmanMohammadi]

@krammnic sounds good! I think there's quite a few techniques out there, so it might be worth prioritizing some of the more commonly used.

[krammnic]

@SalmanMohammadi DPO, KTO, ORPO, SimPO, SPO, SMPO IPO, VinePPO*, LiPO*, R-DPO. Here are some of that I remember which has shown SOTA result somewhere/or pretty commonly used

[bjb19]

Let me take a look at it this weekend and see what a good division might look like. If it is fairly different I can propose a refactor and then start coding it if we are in agreement

[bjb19]

DPO already exists right? So that just leaves us with the other items

[krammnic]

@bjb19 Let me create general issue with possible methods and then we will revise it + set status about DPO, IPO, SPO, etc.

[krammnic]

But another question is should we actually care about other methods that works better in cases where SimPO is limited? Maybe having 2 classic methods + 1 SOTA is actually better strategy if we want to make things less complicated.

[krammnic]

But another question is should we actually care about other methods that works better in cases where SimPO is limited? Maybe having 2 classic methods + 1 SOTA is actually better strategy if we want to make things less complicated.

@SalmanMohammadi What's your opinion about it?

[krammnic]

But another question is should we actually care about other methods that works better in cases where SimPO is limited? Maybe having 2 classic methods + 1 SOTA is actually better strategy if we want to make things less complicated.

@SalmanMohammadi What's your opinion about it?

Need some opinion about these, in another case we can't proceed further

[SalmanMohammadi]

@krammnic there's a couple things on my mind here.

I definitely think it's worth trying to make sure the techniques we're implementing (or raising issues for) will actually be used. Some metrics I might use here is e.g. citations, whether reference implementations exist, models on huggingface trained using X technique.

Generally speaking I don't think we're short of interesting techniques to implement and which we need to prioritize, but rather bandwith to implement these techniques. If we have a wonderful contributor who's interested in implementing a technique I'm mostly up for supporting the implementation. I think it'd be worth continuing a more general discussion on the kinds of RLHF techniques that we could be interested in, in a seperate issue - @krammnic would you be up for creating this issue? All the information you've shared here about the different techniques that are out there would be great to see there, and we can better track there.

I see a bunch of models on hugging face trained with KTO right now, so I'm generally in favour of seeing an implementation in torchtune - if @bjb19 is still interested in contributing this I'd love to keep the discussion going here on how we can make that happen. Did you happen to have a look at this in a bit more detail @bjb19? If not I'm happy to try share what I think would be a good approach to take.