Fix DPO/KTO use_liger_kernel under DeepSpeed ZeRO-3

[qgallouedec]

use_liger_kernel=True + DeepSpeed ZeRO-3 in DPOTrainer crashes. This mirrors the GRPO bug fixed in #5891.

• ✅ liger + zero2
• ✅ no-liger + zero3
• ❌ liger + zero3 → crash

Closes #4544.

Why

The Liger fused DPO loss reads lm_head.weight directly instead of going through model.forward(). Under ZeRO-3 this breaks in two places:

1. Forward. ZeRO-3 shards every parameter to numel 0 and only re-materializes it inside its owning module's forward hook. Since Liger accesses lm_head.weight by attribute (never calling the lm_head module), the gather hook never fires and the fused matmul gets an empty shard → size mismatch, got input (s2x2048), vec (0).
2. Backward. Even once the weight is gathered for the forward, the fused loss produces a dense weight gradient. If the fused loss doesn't run inside the DeepSpeed engine's forward(), ZeRO-3's pre-forward hooks are never armed, so the backward gradient reduction isn't set up and grad accumulation fails → size of tensor a (0) must match the size of tensor b (H).

Fix

Apply the same two-part treatment GRPO already uses (#5891):

1. maybe_gather_lm_head_ctx (new shared helper in trainer/utils.py) all-gathers the lm_head weight/bias for the duration of the fused call. No-op when not on ZeRO-3, or when the weight is already gathered (tied embeddings). Deduplicates by identity so a shared policy/reference head (PEFT with no separate ref_model) isn't passed to GatheredParameters twice.
2. Run _compute_loss_liger through _ForwardRedirection so the fused loss executes inside the engine's forward(), arming the ZeRO-3 pre-forward hooks. The redirection is gated to ZeRO-3 / FSDP: on plain DDP it would wrap Liger's torch.func call inside DDP.forward() and error, so there we call _compute_loss_liger directly on the unwrapped model.

Also routes GRPO's inline gather through the shared maybe_gather_lm_head_ctx (was a duplicated block), per the repo's consistency guideline.

Verified (2×H100)

test_dpo_liger (new, mirrors test_grpo_liger):

backend	result
ddp	✅ pass
zero2	✅ pass
zero3	✅ pass (the fix)
fsdp2	⚠️ xfail — see below

test_grpo_liger[zero3] still passes after routing GRPO through the shared helper.

Reproducer

Before this PR, the following crashes under ZeRO-3; after it, it trains.

accelerate launch --config_file tests/distributed/data/accelerate_configs/zero3.yaml trl/scripts/dpo.py \
    --model_name_or_path trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 \
    --dataset_name trl-internal-testing/zen --dataset_config standard_preference \
    --use_liger_kernel --output_dir /tmp/dpo-liger-zero3

Known limitations / follow-ups

• FSDP2 is xfailed: _compute_loss_liger runs the backbone directly on the unwrapped model, which is incompatible with FSDP2's DTensor-sharded params (mixed torch.Tensor and DTensor). This never worked and is left as a follow-up.
• KTO has the same ZeRO-3 gap, but its Liger path runs an extra full model(...).logits forward in _compute_kl_logps that doesn't cooperate with the redirection. Not fixed here; separate follow-up.

---

Note

Medium Risk
Changes distributed training loss paths (ZeRO-3/FSDP) for DPO/KTO/GRPO/SFT Liger and chunked CE; incorrect gather/redirection could cause silent wrong grads or crashes, but scope is narrow to fused-lm_head paths with new integration tests.

Overview
Fixes use_liger_kernel training under DeepSpeed ZeRO-3 (and aligns FSDP) for DPO and KTO by reusing the same pattern as GRPO: fused losses read lm_head.weight directly, so ZeRO-3 never gathers shards and backward grad reduction breaks.

Shared maybe_gather_lm_head_ctx in trainer/utils.py wraps ZeRO-3 GatheredParameters for lm_head weight/bias (no-op on DDP or when already gathered; dedupes tied/shared params). SFT drops its local copy and GRPO switches to this helper.

DPO/KTO wrap the Liger loss in that gather context, set _ForwardRedirection, and on ZeRO-3/FSDP route compute_loss through redirection so the fused path runs inside the engine forward. Plain DDP still calls _compute_loss_liger on the unwrapped model.

KTO additionally avoids a full model(...).logits pass in _compute_kl_logps when Liger is on (backbone + manual matmul with gather) so ZeRO-3 doesn’t double-touch lm_head.

Adds distributed test_dpo_liger and test_kto_liger; FSDP2 stays xfail (DTensor vs direct backbone/lm_head access).

^{Reviewed by Cursor Bugbot for commit a5b576b. 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.

[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.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 99af755e7b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[kashif]

thanks! useful to have as a util for other trainers, since I was doing a targeted fix