Skip to content

SDFT: add DOPD advantage-gap token routing - #7656

Open
ucalyptus wants to merge 2 commits into
huggingface:mainfrom
ucalyptus:sdft-dopd-routing
Open

ucalyptus wants to merge 2 commits into
huggingface:mainfrom
ucalyptus:sdft-dopd-routing

Conversation

@ucalyptus

@ucalyptus ucalyptus commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Split out of #6237.

Adds distillation_mode="dopd" to SDFTTrainer, following DOPD (Dual On-Policy Distillation, https://huggingface.co/papers/2606.30626). Each completion token routes into one of four regimes, chosen from the advantage gap and each side's confidence:

  • low gap, either side confident: light top-k reverse KL toward the teacher
  • low gap, both unsure: weak self-regularization with a stop-gradient anchor
  • high gap, teacher confident: full-vocabulary JSD toward the teacher
  • high gap, student confident: light stop-gradient privileged-student consistency nudge

Tokens where the gap is high but neither side is confident fall back to the weakest term, matching the paper's intent that ambiguous tokens get the smallest update.

Routing measures the gap between the privileged teacher and the privileged student, which is the student's own no-grad forward on the same privileged context. Using the bare-prompt student instead mixes the capability gap with the information-asymmetry gap and inverts regime selection; the code comments call this out where it matters.

teacher_model_kind="live" is rejected with this mode. The teacher and the privileged student would hold the same weights on the same input, the gap would collapse to zero, and the two high-gap regimes would be unreachable. use_teacher_server=True is rejected for the same reason as full_logits: the server returns top-k logprobs only, and the routing needs full logits for the confidence and gap terms.

Note on topk_support

compute_topk_self_distillation_loss gains a topk_support parameter here, because the DOPD loss mixes a teacher-support regime with two student-support regimes. #7655 adds the same parameter, along with the config field, the teacher-server alignment, and tests for it. If #7655 merges first, I will rebase and this hunk disappears, leaving only the routing code.

Tests

CPU, hand-built logits, no model or GPU:

pytest tests/experimental/test_dopd_routing.py tests/experimental/test_self_distillation_trainer_behavior.py -q
20 passed

test_dopd_routing.py checks that each row lands in its expected regime, that the four regime masks are mutually exclusive and exhaustive, that regimes 2 and 4 anchor on the privileged student rather than the bare student (a bare anchor makes the term identically zero, in value and in gradient), and that every regime backprops a nonzero gradient into the student.


Note

Medium Risk
Changes the SDFT training loss path and adds a multi-regime objective with new hyperparameters; default distillation_mode stays topk_logits, but DOPD needs an extra privileged forward per step and rejects some teacher/vLLM setups.

Overview
Adds DOPD (Dual On-Policy Distillation) to experimental SDFT via distillation_mode="dopd", so each completion token gets a different distillation objective instead of one global mode.

Routing uses the advantage gap between the privileged teacher and a privileged-student forward (same weights, teacher prompt, no_grad) plus per-side confidence. Four regimes: light teacher top-k KL, weak self-reg against stop-gradient privileged student, full-vocab JSD, or light student-consistency; ambiguous high-gap tokens fall back to self-reg.

Config adds distillation_dopd_* thresholds/weights; teacher_model_kind="live" and use_teacher_server=True are rejected for DOPD (gap collapse / no full logits). compute_topk_self_distillation_loss gains optional topk_support ("student" vs "teacher") for mixed DOPD regimes.

Tests (test_dopd_routing.py) cover regime selection, privileged-vs-bare gap, gradients, and config validation. Docs updated in sdft_trainer.md and paper_index.md.

Reviewed by Cursor Bugbot for commit c30e622. Bugbot is set up for automated code reviews on this repo. Configure here.

@ucalyptus
ucalyptus force-pushed the sdft-dopd-routing branch 2 times, most recently from e3088a7 to 495c84a Compare October 13, 2026 00:00
@ucalyptus
ucalyptus marked this pull request as ready for review October 13, 2026 00:03
Adds `distillation_mode="dopd"`, following DOPD (Dual On-Policy Distillation,
https://huggingface.co/papers/2606.30626). Each completion token routes into one
of four regimes by advantage gap and confidence: light top-k reverse KL toward
the teacher, weak stop-gradient self-regularization, full-vocabulary JSD toward
the teacher, or a light stop-gradient privileged-student consistency nudge.

Routing measures the gap between the privileged teacher and the privileged
student (the student's no-grad forward on the same privileged context); using
the bare-prompt student instead conflates the capability gap with the
information gap and inverts regime selection.

`teacher_model_kind="live"` is rejected in this mode, because both policies
would hold the same weights and the high-gap regimes would be unreachable.
ucalyptus added a commit to ucalyptus/trl that referenced this pull request Oct 13, 2026
The combined branch: DOPD advantage-gap routing, the distillation plumbing
(kl_clip, teacher top-k support, teacher chat-template kwargs with untruncated
teacher tokenization, QLoRA, EMA teacher adapter pre-registration, parallelism
rejection), and the cross-rank metric averaging fix, with the config validation
and tests for each.

Split into huggingface#7654 (metric), huggingface#7655 (plumbing), and huggingface#7656 (DOPD) for review. The
PEFT + ZeRO-3 guard that this branch also carried was dropped: it rejected a
combination that is known to work, and huggingface#6829 owns that fix.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant