[DistillationTrainer refactor] Heavy-clean the Liger path: share extraction with the chunked path

[qgallouedec]

Item 34 (Group E) — stacked on PR 33. Part of #6449.

Keeps the Liger fused-JSD fast path (use_liger_kernel=True) but heavy-cleans it: the two loss paths now share the same hidden-state extraction and differ only in the final loss call. Net −74 lines.

compute_loss always redirects _compute_loss through _forward_redirection (from item 33). _compute_loss extracts student hidden states (inside the student wrapper) + teacher hidden states (via the teacher's own redirection), then:

• if self.use_liger_loss: → self.liger_loss(...) over the flattened hidden states, with true_labels built from completion_mask (hard-loss weight is 0, so only the ignore-mask matters; no shift — _get_last_hidden_state is already completion-aligned), wrapped in the shared maybe_gather_lm_head_ctx for ZeRO-3, then the num_items global rescale.
• else → _chunked_divergence_loss(...).

Deleted (all superseded by the shared extraction + shared gather util):

• _liger_student_forward, _compute_liger_loss, _get_liger_zero3_lm_head_gather_ctx (and its wrapping in training_step/prediction_step);
• the hasattr/getattr decoder/bias probing;
• the unused nullcontext import.

No use_liger_kernel raise (reversing the earlier "delete Liger" plan — SFT parity, and the path is verified working). The config field stays.

Test: test_distillation_trainer_with_liger (train-smoke) replaced with test_liger_loss_agrees_with_chunked — asserts the Liger loss equals the chunked loss on a fixed batch (rtol=1e-3), toggling use_liger_loss on one un-patched model so only the loss kernel differs. Suite green (52).

Known limitation (pre-existing, unchanged): the Liger kernel does not apply logit_scale/final_logit_softcapping, so on Gemma/Cohere it diverges slightly from the chunked path; the chunked default is exact. Worth a follow-up if Liger + those families matter.

---

Note

Medium Risk
Changes core distillation loss routing and distributed (DDP/FSDP/ZeRO-3) behavior for the Liger path, though parity is now explicitly tested against the chunked path.

Overview
DistillationTrainer refactors the optional Liger fused-JSD path so it uses the same student/teacher hidden-state extraction and _forward_redirection wiring as the default chunked loss; only the final loss kernel differs.

compute_loss no longer branches early for Liger—it always routes _compute_loss through the DDP/FSDP wrapper. Inside _compute_loss, when use_liger_loss is set, it calls LigerFusedLinearJSDLoss on the shared hidden states with completion_mask-derived true_labels, maybe_gather_lm_head_ctx for ZeRO-3 lm_heads, and the same num_items_in_batch rescale. The separate _liger_student_forward, _compute_liger_loss, and _get_liger_zero3_lm_head_gather_ctx helpers are removed, along with ZeRO-3 gather wrappers in training_step and prediction_step.

The Liger integration test is replaced with test_liger_loss_agrees_with_chunked, which asserts Liger and chunked losses match on a fixed batch when only use_liger_loss toggles.

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eb744e799f

ℹ️ 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".

[chatgpt-codex-connector]

Detect logits_to_keep support on the teacher

When use_liger_kernel=True with different student and teacher architectures, this new common routing sends the Liger path through _get_last_hidden_state for the teacher. That helper decides whether to pass logits_to_keep using self.model_kwarg_keys, which was derived solely from the student signature, so a student that supports the argument paired with a same-vocabulary teacher backbone that rejects it now raises TypeError; the previous Liger implementation called the teacher decoder without this argument. Determine support separately for each model, or pass the appropriate capability into the shared extractor.

Useful? React with 👍 / 👎.