SDFT/SDPO: distillation plumbing fixes

[ucalyptus]

Split out of #6237, which grew to cover several unrelated changes. This is the second piece, after #7654 (metric averaging); the two do not overlap.

Fixes and small features in the SDFT/SDPO distillation path. They are separate concerns that happen to touch the same config and loss plumbing, so they are grouped here rather than in one PR each. Happy to split further if you want them apart.

• distillation_kl_clip caps the summed per-token divergence before averaging, so outlier tokens cannot dominate the signal. Supported for full_logits and topk_logits.
• distillation_topk_support lets topk_logits project onto the teacher's top-k logits. With use_teacher_server=True the field is aligned to "teacher", because the vLLM server only returns the teacher's top-k.
• teacher_chat_template_kwargs forwards extra apply_chat_template kwargs when the teacher prompt is built, for cases like a thinking teacher paired with a non-thinking student. The teacher is also scored on the untruncated prompt now, since the problem leads the teacher template and left-truncation would drop it. Generation used the truncated prompt before, so teacher logits and completions could disagree.
• quantization_config accepts a BitsAndBytesConfig for QLoRA, and trainable adapter weights are cast to bf16 per the QLoRA paper (paragraph 3 of https://huggingface.co/papers/2305.14314).
• The EMA "teacher" adapter is registered in __init__ instead of at the first train() step, so it exists before accelerate or DeepSpeed wraps the model. ZeRO-3 registers each module once at initialization and cannot adopt later additions. PEFTAdapterEMACallback no longer re-adds an adapter that is already registered; PEFT rejects a duplicate name.
• parallelism_config values that the trainer cannot support (context or sequence parallel) are rejected at construction.
• generation_kwargs and chat_template_kwargs are listed as valid dict CLI fields, and both may be passed as a string.

SDFT and SDPO receive the same changes. The two trainers are self-contained by design, so the code is duplicated rather than shared.

Tests

tests/experimental/test_distillation_plumbing.py is new: eleven CPU-only tests. Six cover the loss plumbing on hand-built logits (topk_support picks different token sets for student and teacher support and rejects an invalid value; kl_clip bounds the summed divergence, None disables it, and it is rejected for modes without a full logit distribution; teacher-server alignment). Five more cover the config and init paths: SDPO parity for both new fields, and an EMA pre-registration test per trainer that asserts the "teacher" adapter exists straight after construction, without train().

CPU, trl-internal-testing/tiny-Qwen2ForCausalLM-2.5, one file at a time:

pytest tests/experimental/test_distillation_plumbing.py -q              11 passed
pytest tests/experimental/test_sdft_trainer.py -q                      18 passed
pytest tests/experimental/test_sdpo_trainer.py -q                      21 passed
pytest tests/experimental/test_self_distillation_trainer_behavior.py -q 12 passed

The new config fields are also exercised by the behavior test's trainer construction, which now passes distillation_topk_support and distillation_kl_clip.

Not covered by a test: the QLoRA cast's DoRA exclusion, which needs a quantized model.

---

Note

Medium Risk
Changes affect core distillation loss math, teacher prompt tokenization, and PEFT/DeepSpeed initialization order; training dynamics can shift when new options are used, though defaults largely preserve prior behavior.

Overview
This PR extends SDFT and SDPO distillation plumbing with configurable top-k support, optional KL capping, and several trainer/init fixes mirrored in both code paths.

Loss and config: distillation_topk_support chooses whether top-k logit distillation uses the student or teacher vocabulary support (default "student"); with use_teacher_server and topk_logits, config is auto-aligned to "teacher". distillation_kl_clip caps the summed per-position divergence (after vocab sum) for full_logits and topk_logits, and is rejected for sampled_token. Shared loss_utils in both packages wire these through compute_divergence and the top-k/full logit objectives.

SDFT-specific: teacher_chat_template_kwargs and untruncated teacher prompt tokenization keep teacher scoring/generation consistent when the problem must not be left-truncated away.

Trainer init: Optional quantization_config for QLoRA; trainable LoRA weights cast to bf16 on quantized models (DoRA magnitude excluded). EMA "teacher" PEFT adapter is registered and zeroed in __init__ for ZeRO-3 compatibility; PEFTAdapterEMACallback skips re-adding an existing adapter. Unsupported context/sequence parallelism_config is rejected at config time. Dict CLI fields expanded for generation/chat template kwargs (including string forms).

Tests: New test_distillation_plumbing.py (CPU, hand-built logits) plus behavior test defaults updated for the new config fields.

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

[cursor]

Stale Bugbot comment from a previous run.

[ucalyptus]

All three are fixed in 2c283a0.

SDPO skips EMA adapter pre-registration (high). SDPOTrainer.__init__ now registers the "teacher" adapter and zeroes it before accelerate or DeepSpeed wraps the model, mirroring the block SDFT already had. SDPO does default to teacher_model_kind="ema", so this was a real gap, and the teacher_sync guard did assume the trainer had already registered the adapter.

SDPO omits new distillation config wiring (medium). SDPOConfig gains distillation_topk_support and distillation_kl_clip with the same defaults, docstrings, mode validation, and teacher-server alignment as SDFTConfig, and SDPOTrainer passes both at the top-k call site, kl_clip at the full_logits call site, and kl_clip on the teacher-server path. The SDPO loss helpers are no longer dead parameters.

DoRA magnitude vector (medium). The QLoRA cast in both trainers now iterates named_parameters() and skips lora_magnitude_vector, with the same comment the main trainers carry (see #7268).

test_distillation_plumbing.py grows from 6 to 11 tests: SDPO config parity (defaults, the kl_clip mode rejection, teacher-server alignment) and an EMA pre-registration test per trainer that asserts "teacher" in trainer.model.peft_config right after construction, without calling train(). 60 tests pass across the four experimental files (plumbing 11, sdft 17, sdpo 20, behavior 12).

One gap worth stating: the DoRA exclusion has no test, because it needs a quantized model, which the experimental suite does not build. The change is a line-for-line match with the main trainers instead.

[DimensionSTP]

@ucalyptus, #6829 is now updated at c99446a4 following qgallouedec's requested callback design. Both SDPO and SDFT create the teacher adapter in PEFTAdapterEMACallback.__init__, without a private initialization call from the Trainer. ZeRO-3 EMA synchronization also gathers the student and teacher adapter parameters without an empty-shard condition and writes the update on rank zero.

This overlaps the Trainer-side EMA pre-registration blocks here. I propose landing the callback initialization and synchronization repair through #6829, while keeping your config, loss, teacher-server and QLoRA changes in #7655. Rebasing on that fix would let you remove the overlapping Trainer-side initialization calls without losing those independent changes.

Both official Trainer paths pass two-rank ZeRO-3 training with checkpointing off and on, including exact teacher/student equality after two steps at update rate 1. Could you align the overlapping initialization changes with #6829 so we avoid introducing two initialization paths for the same adapter?

[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 1a92fbe. Configure here.}