Support Chunked DPO Loss Kernel

[austin362667]

Summary

Add support for a fused, torch-compiled, and chunked DPO (Direct Preference Optimization) loss kernel, as requested in #371.
This implementation is largely based on the excellent work done on ORPO (#362) by @shivam15s.

DPO Loss Formulation

In a reference setting:

$$r_\theta(x,y_c) - r_\theta(x,y_r) = \log(\pi_\theta(y_c|x)) - \log(\pi_\theta(y_r|x))$$

$$-\log(\sigma((\log(\pi_\theta(y_c|x)) - \log(\pi_\theta(y_r|x)) - \log(\pi_{\theta_{\text{ref}}}(y_c|x)) + \log(\pi_{\theta_{\text{ref}}}(y_r|x)))/\beta))$$

Corresponds to:

# Policy model log probabilities
policy_chosen_logps = log_probs(policy_chosen_logits)
policy_rejected_logps = log_probs(policy_rejected_logits)

# Reference model log probabilities
ref_chosen_logps = log_probs(ref_chosen_logits)
ref_rejected_logps = log_probs(ref_rejected_logits)

# Compute advantages
chosen_advantages = policy_chosen_logps - ref_chosen_logps
rejected_advantages = policy_rejected_logps - ref_rejected_logps

# policy_chosen_logps - ref_chosen_logps - policy_rejected_logps + ref_rejected_logps
logits_diff = (chosen_advantages - rejected_advantages) * beta

# DPO loss
losses = -F.logsigmoid(logits_diff)

Testing Done

• Hardware Type: NVIDIA L40S (48G)
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[austin362667]

Should I use HF DPO impl here in benchmarking for function reusability purpose? Or write another naive impl in pure torch?

[lancerts]

HF DPO should be fine

[austin362667]

Should I move this HF impl to file test_dpo_loss.py?

[Tcc0403]

yes, since HF impl is only for testing purpose

[lancerts]

can we modify logits_diff = (chosen_logps - rejected_logps) / beta
to
logits_diff = (chosen_logps - rejected_logps) * beta to align with the convention in paper as well as the trl implementation here

[pramodith]

Just a FYI, I think we should wait until @shivam15s pushes a generic/inheritable class that handles all the chunking and other repetitive logic common to different loss functions, before pushing new loss functions.

[shivam15s]

Great work @austin362667 ! The additional summing of NLL loss is going to be useful for IRPO loss as well :). I'll be creating a simple base class which adds the boilerplate code (backward/torch compile logic) that you can inherit from, as @pramodith mentioned

[austin362667]

Issue addressed. Thanks @Tcc0403 @lancerts @pramodith @shivam15s and @ByronHsu for review!

[Tcc0403]

I think we should make chunked_loss functions nn.Module (like flce and fljsd) for users? same for orpo? cc @shivam15s @ByronHsu

[ByronHsu]

@Tcc0403 that is the plan!