Move DPO to chunked log probabilities

[kashif]

What does this PR do?

Moves DPO onto the shared chunked log-probability path while keeping the trainer loss math in one place. This removes the duplicated fused DPO loss and lifts its restrictions on mixed losses, f-divergences, and precomputed reference log-probs.

Part of #7063.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that is the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that is the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

---

Note

High Risk
Changes the core DPO forward/backward path (memory, numerics, distributed hooks) and deletes the old fused implementation; regressions would affect all use_liger_kernel DPO training.

Overview
Replaces the separate fused DPO loss with a shared chunked log-probability path so use_liger_kernel=True no longer depends on FusedLinearDPOLoss / liger-kernel for DPO. Policy and reference log-probs are computed via backbone forward + _ChunkedLogProbFunction (vocab-chunked lm_head), then the existing multi-loss / f-divergence logic in _compute_loss runs unchanged.

Removes trl/losses/dpo_loss.py, fused_linear_preference.py, and the large tests/test_dpo_loss.py parity suite. Drops init-time blocks on multiple loss_type, non-default f_divergence_type, and precompute_ref_log_probs; only use_weighting=True (and compute_metrics, MoE aux, PEFT edge cases) stay incompatible. Adds FSDP2 support by gathering sharded lm_head DTensors before projection.

Docs and distributed tests are updated (FSDP2 no longer xfail for Liger DPO). Trainer tests now compare chunked vs full-logit loss/gradients across loss types and f-divergences.

^{Reviewed by Cursor Bugbot for commit c2196d5. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Benchmarked it (single H100, Qwen3-0.6B, bs=4, seq 512)

config	median step	peak
main, fused, inline ref	0.2243 s (n=4)	7.05 GB
this PR, inline ref	0.2457 s (n=6)	7.05 GB
this PR, precompute_ref_log_probs=True	0.1971 s (n=2)	4.83 GB

+9.5% per step at equal config, -12.1% with precomputed ref logps, which main rejects outright. Memory unchanged in the first case, 31% lower in the second.

So nothing unexpected!

A few things:

• For the record, I found that compute_loss redirects on is_zero3 or self.is_fsdp_enabled, so plain DDP falls through to _compute_loss(unwrapped_model, ...) and prepare_for_backward() never fires. That is inherited from main so leaving it is right. I'll open an issue to fix it across trainers.

• Docs need updating, dpo_trainer.md has a "Compatibility and constraints" section listing, under use_liger_kernel=True: only a single loss_type, no compute_metrics, no precompute_ref_log_probs. This PR lifts the first and the third and adds a new one (use_weighting=True), so three lines there are now wrong.

• Same section is where the f-divergence lift belongs, it was only ever enforced in code.

[qgallouedec]

Related to #7245: DPO silently skips DDP gradient synchronization with use_liger_kernel=True, and this PR is what makes it fixable.

compute_loss redirects only on is_zero3 or self.is_fsdp_enabled, so under plain DDP the loss runs on the unwrapped model, DistributedDataParallel.forward() never runs, its reducer is never armed, and gradients are never all-reduced. Measured on 2 GPUs with different data per rank: 310/310 parameters differ across ranks, on main and on this branch alike.

On main the one-line fix does not work. Adding or model is not unwrapped_model makes every rank raise NotImplementedError: Cannot access storage of TensorWrapper, because FusedLinearDPOFunction runs torch.func.grad_and_value and that cannot run inside DistributedDataParallel.forward().

On this branch the same change gives 0/310. _ChunkedLogProbFunction is a plain autograd.Function, so the redirect is fine.

So the DPO half of #7245 belongs here, as one line in compute_loss:

if is_zero3 or self.is_fsdp_enabled or model is not unwrapped_model:

KTO is going separately, it is already on the chunked path.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c2196d5. Configure here.}

[cursor]

DDP redirect not copied to KTO

Medium Severity

compute_loss now redirects through the wrapper whenever model is not unwrapped_model, so DDP arms its reducer on the chunked path. The same duplicated block in KTOTrainer.compute_loss still redirects only for ZeRO-3 and FSDP. Copied trainer logic in this repo has to stay aligned, and KTO still skips DDP gradient synchronization with use_liger_kernel=True.

 

^{Triggered by project rule: ../.ai/AGENTS.md}

^{Reviewed by Cursor Bugbot for commit c2196d5. Configure here.}

[qgallouedec]

Thanks!, both points addressed. Three things left:

• The guard has to match KTO on main (Redirect the KTO liger loss through the distributed wrapper under DDP #7247): if self.is_fsdp_enabled or model is not unwrapped_model:, and the deepspeed_plugin / is_zero3 lines go. Same block, same text.
• Merge main, the branch predates Redirect the KTO liger loss through the distributed wrapper under DDP #7247.
• Min-versions CI is red on all five gradient tests, and it's not noise: the mismatch sits on a few rows of one 1.2M-element parameter. Reduction-order differences don't look like that. Worth finding what the old torch does differently there before touching the tolerance.