DPO Loss refactor

[SalmanMohammadi]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• refactor

I was initially looking at this refactor to add SimPO into the DPO recipe, but got slightly carried away. I think we might need a separate SimPO recipe down the line anyway, since DPO-style losses use a reference model, whereas SimPO doesn't.

This PR refactors the DPO loss module into separate classes for each of the loss types it supported. Each separate loss is now documented with a reference to its corresponding paper, a little intuition about how it works, and comes with a corresponding unit test.

Test plan

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1197

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 55ccd92 with merge base 0ea274f ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[RdoubleA]

Excellent refactor. I honestly don't have any concerns here.

As a DPO noob, do you mind adding some details in the docstrings about the training setup for each of these losses? Like what is the policy/reference model, and the core value prop of DPO compared to PPO is that it does not require a reward model.

Also, are we completely removing the kto pair loss?

[SalmanMohammadi]

Also, are we completely removing the kto pair loss?

Yeah, sorry, I forgot to point this out. The reference implementation in TRL seems to differ enough (it uses a separate trainer with a different dataset format), and I couldn't see tests for the KTO loss in the original DPO PR. I personally would re-consider this when we have tests against a reference implementation.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 68.66%. Comparing base (f4f83b2) to head (55ccd92).
Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1197      +/-   ##
==========================================
+ Coverage   68.63%   68.66%   +0.02%     
==========================================
  Files         214      215       +1     
  Lines        9687     9734      +47     
==========================================
+ Hits         6649     6684      +35     
- Misses       3038     3050      +12     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.