[DistillationTrainer refactor] Wire compute_loss to the chunked JSD path; delete the full-logit loss

[qgallouedec]

Item 33 (Group E) — stacked on PR 32. Part of #6449. The loss handover.

Switches compute_loss from the full-logit generalized_jsd_loss to the chunked path added at item 32, then deletes the old loss. The teacher's dense distribution is now matched without ever materializing (B, C, V) logits.

New compute_loss (non-Liger branch):

• Student backbone via _get_last_hidden_state routed through _forward_redirection (so DDP.forward / prepare_for_backward fire); teacher backbone via _get_last_hidden_state under no_grad.
• _chunked_divergence_loss(...) over the two hidden states + each model's lm_head, with per-model logit_scale/final_logit_softcapping sourced from config (getattr, as SFT does).
• _forward_redirection moved to unconditional ctor state (was Liger-only; the chunked path needs it too).

Deleted (orphaned): generalized_jsd_loss, _jsd_divergence, _reduce_divergence_loss, _get_teacher_logits. Siblings (GKD/GOLD/IW-OPD/server) keep their own copies. Liger path untouched (removed at item 34).

Handover proof: test_compute_loss_matches_full_logit_reference asserts the chunked compute_loss equals a from-scratch full-logit JSD over the completion positions (assert_close; exact at the tiny model's fp32). The item-32 unit tests continue to pin the chunk math.

Tests: deleted TestGeneralizedJSDLoss, TestGeneralizedJSDLossIsPinned, _reference_generalized_jsd, and test_num_items_in_batch_counts_the_tokens_trained_on (its _reduce_divergence_loss recording closure is gone). test_loss_normalizes_by_num_items_in_batch kept (routes through the chunked path now). Suite green.

Subsumes item 35 (temperature). The chunked loss has no temperature term (SFT parity), so compute_loss no longer divides logits by self.temperature — temperature is now sampling-only, which was item 35's goal. At the default temperature=1.0 this is exactly equivalent to the old loss (the handover proof holds); it changes loss values only for temperature != 1.0. Item 35 reduces to a docstring/cleanup (or drops).

---

Note

Medium Risk
This changes the core training loss and distributed forward path for all non-Liger distillation runs; behavior is guarded by handover and normalization tests but still touches gradient-critical code.

Overview
DistillationTrainer no longer materializes full (batch, completion, vocab) logits for the default loss. compute_loss runs hidden states from student and teacher backbones (via _get_last_hidden_state) into _chunked_divergence_loss, with _forward_redirection on both models so DDP/FSDP prepare_for_backward and sharded lm_head behavior stay correct. _forward_redirection is always set in the constructor, not only when Liger is enabled.

The legacy stack is removed from this trainer: generalized_jsd_loss, _jsd_divergence, _reduce_divergence_loss, and _get_teacher_logits. Chunked _chunk / _chunked_divergence_loss now accept temperature (applied after per-model scale/softcapping) and compute_loss passes self.temperature.

Tests drop the old full-logit generalized_jsd_loss suites and add test_applies_temperature on the chunked helper plus test_compute_loss_matches_full_logit_reference as a handover check at default temperature; test_loss_normalizes_by_num_items_in_batch is updated for the chunked path.

^{Reviewed by Cursor Bugbot for commit 71d4c01. 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: 92eee55b4b

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

[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 80f3006. Configure here.}