[DistillationTrainer refactor] Accept a prompt column alongside messages#6461
Merged
Conversation
The trainer is about to go through a long refactor (top-k support removal,
then a switch to a chunked loss). The implementation is expected to change;
the value it computes is not.
Add three pinning tests:
- `generalized_jsd_loss` against a naive reference written straight from the
definition. It builds the mixture in probability space rather than reusing
`F.kl_div`'s inverted argument order and the `logsumexp` trick, so it can
actually catch an argument-order or mixture-weight regression.
- parity with `GKDTrainer.generalized_jsd_loss`, which implements the same
objective. This is the cross-trainer contract.
- the temperature path, which is scheduled to become sampling-only. Pinning it
makes that change a visible diff rather than a silent drift.
All are seeded and cover beta in {0, 0.5, 1}, with and without a label mask.
The existing fixture uses an unseeded `torch.randn`, which cannot pin anything.
`lmbda` selects between training on the dataset's own completions (`lmbda=0.0`) and on completions the student generates (`lmbda=1.0`). The trainer is scheduled to become always-on-policy, so pin both modes now: the later removal of the off-policy path should be a deliberate deletion with a visibly failing test, not a silent behaviour change. Asserts the loss is recorded and that every parameter actually moved, using `torch.equal` rather than a tolerance so a stalled update cannot pass. Uses a higher learning rate than the default because gradients are tiny on the test model and the default rate can stall the update, which would make the assertion vacuous.
The divergence loss is reduced as `sum / num_items_in_batch`, so that count is the loss denominator. It is currently wrong. `transformers.Trainer` computes it from the raw dataloader batches, before `_prepare_inputs` runs. Two things follow. `_RepeatBatchDataLoader` yields the same generation batch once per accumulation step, so the count is `gradient_accumulation_steps` times too large. And it is derived from the dataset labels, even on steps whose completions are replaced by generated ones. The effect is a silently down-scaled loss and gradient. Add an xfail test that records what the trainer is actually handed during `train()` and compares it with the tokens the loss is summed over. The existing coverage passes `num_items_in_batch` explicitly, so it exercises the reduction rather than the value, and cannot catch this. Un-xfail when the generation buffer moves to the GRPO-style `_prepare_inputs` and the denominator is derived from the completions actually trained on.
IW-OPD (Importance-Weighted On-Policy Distillation, 2606.22600) is a teacher-guided policy-gradient method: sampled-token advantage `A_t = sg[log π_T − log π_0]`, a prefix-decayed importance weight, and a clipped-surrogate PPO inner loop. It fits neither the stable full-vocabulary `DistillationTrainer` (supervised KD, no advantages) nor `GRPOTrainer` (which has no teacher). So rather than delete it, move it to its own home. - Fork the pre-refactor trainer/config into `trl/experimental/iw_opd/` as `IWOPDTrainer` / `IWOPDConfig` — a self-contained snapshot that keeps IW-OPD (and, incidentally, the other pre-refactor features) working. Its docstring states plainly that it is a frozen, unmaintained snapshot. - Remove IW-OPD from the base `DistillationTrainer` (loss, config fields, rollout-logprob plumbing, the `logprobs=0` vLLM request, and the tests). - Repoint the IW-OPD paper-index entry to `experimental.iw_opd` (the GKD and general on-policy-distillation entries stay on `DistillationTrainer`). - Register `IWOPDTrainer` in the telemetry allowlist; the moved tests live in `tests/experimental/test_iw_opd_trainer.py`. This merges what were two PRs (remove IW-OPD + drop its paper entry) into one "move" so nothing is lost. `compute_loss` in the base still drops from 7 dispatch paths to 5.
The server-backed path (fetch per-token teacher logprobs from a vLLM server instead of a local teacher forward) is a distinct use case that only ever runs over a sparse top-k support. Move it out of the base trainer into a dedicated experimental subclass so the base can converge on the local, full-vocabulary objective. New package `trl/experimental/server_distillation/`: - `ServerDistillationTrainer(DistillationTrainer)` overriding `compute_loss` with the server path, plus the moved `_get_teacher_token_logprobs_from_server`, `_compute_server_sparse_top_1_divergence_loss`, `_compute_server_forward_kl_loss` and the module-level `build_teacher_request_inputs`. The shared sparse helpers (`_compute_sparse_top_1_divergence_loss`, `_get_reverse_kl_top_1_tokens`, `_add_tail_bucket`, `_jsd_divergence`, `_reduce_divergence_loss`) are inherited from the base, not copied — they are still used by the base local top-1 path. - `ServerDistillationConfig(DistillationConfig)` adding `teacher_model_server_url` and carrying the server-only validations, now unconditional. Base trainer: - `__init__` drops the server branch; a `None` teacher just leaves `self.teacher_model = None`. - `compute_loss` drops the server dispatch; the local path runs unconditionally. - `_get_teacher_logits` drops the server NotImplementedError branch. - config drops `use_teacher_server` / `teacher_model_server_url` and their checks. Server tests move to `tests/experimental/test_server_distillation_trainer.py`. The server-side infrastructure (`/get_sequence_logprobs/`, VLLMClient) is untouched and still used by SDFT/SDPO. `compute_loss` in the base goes from 5 dispatch paths to 2 (local top-1, local full-vocab) plus the Liger path.
With the server path extracted, the base trainer no longer needs the mixed top-1 support path: the local teacher always has full logits, so it can compute the exact generalized JSD/KL directly. - `compute_loss` local branch now always calls `generalized_jsd_loss` (which still honours `loss_top_k` for the top-k approximation); the `beta > 0 and loss_top_k == 1` special case and its `completion_tokens` slice are gone. - `_compute_local_sparse_top_1_divergence_loss` is deleted. - `_compute_sparse_top_1_divergence_loss` and `_get_reverse_kl_top_1_tokens` become server-only, so they move to `ServerDistillationTrainer`. The subclass now defines them directly rather than inheriting them. - `reverse_kl_top_1_mode` only drove those helpers, so it moves from `DistillationConfig` to `ServerDistillationConfig` (field + the "must be one of" validation), and the `self.reverse_kl_top_1_mode` assignment moves from the base `__init__` to the subclass `__init__`. Its stale mentions in the `loss_top_k` docs are trimmed. - The corresponding config test moves to the server test file. Verified on CPU that the server sparse loss is numerically unchanged by the relocation.
The base trainer always has full teacher logits, so it never needs the top-k
approximation — only the extracted server path (which sees just the teacher's
top-k logprobs) does.
- `generalized_jsd_loss` drops the `top_k` / `add_tail` parameters and the whole
top-k support-selection block; it is now the pure full-vocabulary path.
- config `loss_top_k` / `loss_add_tail` move from `DistillationConfig` to
`ServerDistillationConfig`; the `self.loss_top_k` / `self.loss_add_tail`
assignments move from the base `__init__` to the subclass `__init__`; the
base `compute_loss` call drops the two kwargs.
- `_add_tail_bucket` was only reachable from the top-k block and the server
methods, so it moves to the server trainer module. `_jsd_divergence` stays
shared in the base module (the base uses its dense branch, the server its
masked branch).
- `test_loss_normalizes_by_num_items_in_batch` drops its `loss_top_k`
parametrization (base is full-vocab only); the `_add_tail_bucket` import in
the server test now comes from the server module.
Verified on CPU that `generalized_jsd_loss` still matches a naive full-vocab
reference for beta in {0, 0.5, 1} with and without labels, and that the server
sparse loss is numerically unchanged.
The base local-teacher loss compares full next-token distributions, so the only real requirement is a shared vocabulary. Replace the tokenizer-loading / `get_vocab()` comparison + warn + `resize_token_embeddings` dance with a strict `config.vocab_size` equality check that raises (pointing at GOLD for the cross-tokenizer case), mirroring GKD. - `__init__` teacher setup no longer loads the teacher tokenizer, no longer requires `config._name_or_path`, and no longer resizes the teacher embeddings. The check is gated on `teacher_model is not None`, so subclasses that supply the teacher another way (ServerDistillationTrainer) are unaffected. - Remove `_local_teacher_tokenizers_match`, `_raise_if_local_teacher_tokenizer_mismatch`, the `self._local_teacher_tokenizer_matches_student` flag, and the guard call at the top of `compute_loss`. Drop the now-unused `warnings` import. - Add tests for the vocab-size mismatch (Qwen2 student + Llama teacher) and for passing `teacher_model_init_kwargs` alongside an already-instantiated teacher. Note: `teacher_model_name_or_path` remains an unused config field for now; it is removed later when the teacher becomes a constructor-only argument.
The distillation trainer is converging on always-on-policy training (the GRPO-shaped stable target generates every batch). `lmbda`, which mixes on-policy (student-generated) and off-policy (dataset) completions per gradient-accumulation slice, is being removed. Warn with FutureWarning when `lmbda != 1.0`, pointing at `GKDTrainer`, which exposes the same knob for users who need on/off-policy mixing. The default (`lmbda=1.0`, fully on-policy) does not warn. Deprecate-before-remove: the removal lands in a later PR. Both policy modes are already pinned end to end (see the on/off-policy training test), so that removal will be a visible, guarded change.
The GKD (2306.13649) reproduction sets `lmbda=0.5` for the paper's on/off-policy mixing. `DistillationTrainer` is dropping `lmbda` (becoming always on-policy), so move the reproduction snippet to `GKDConfig` / `GKDTrainer`, which keep that knob. Translate `max_completion_length` to GKD's `max_new_tokens`. The prose now frames GKDTrainer as the home for the paper's `lmbda` mixing and DistillationTrainer as the same generalized-JSD objective for the always-on-policy case. Must land before `lmbda` is removed from DistillationConfig, so the index never references a removed argument.
The trainer is now always on-policy: every generation-batch slice is generated by the student, and there is no dataset-completion (off-policy) path. Trainer: - `_fill_buffer` generates for every slice; the per-slice on/off-policy coin flip (`random.random() <= lmbda`, broadcast across ranks) and the off-policy "store the dataset slice directly" branch are gone. - Remove `self.lmbda`, `self._buffered_on_policy_flags`, and the four on/off-policy loss accumulators. `training_step` drops the on/off-policy attribution (its completion-length metrics lose the `on_policy_`/`off_policy_` prefix), and `log` drops the `on_policy_loss` / `off_policy_loss` emission and the all-reduce of the accumulator vector. - Drop the now-unused `random` coin flip, `broadcast_object_list`, and `torch.distributed` / `DistributedType` imports (rich still uses `random`). Config: - Remove `lmbda`, its range check, the deprecation warning added last PR, and the `num_generations > 1 and lmbda < 1.0` warning. Tests: - `_make_args` drops `lmbda=0.0`; the on/off-policy training test becomes a single always-on-policy `test_train_updates_params`; the `lmbda` deprecation tests are removed; the server ragged-grad test drops `lmbda=0.0`. Both policy modes were pinned before this (PR 02 / PR 10), so the off-policy removal is a deliberate, test-visible deletion. Full experimental suite green.
The model card built by IWOPDTrainer attributed runs to the OPD paper (2306.13649) carried over from the fork source. Point _paper at the IW-OPD paper (2606.22600, "On the Position Bias of On-Policy Distillation") so published checkpoints cite the correct work.
loss_top_k was removed from DistillationConfig (it now lives on ServerDistillationConfig), so the local usage guide no longer documents it. The "three key parameters" count drops to two (lmbda, beta).
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 16ec8c3. Configure here.
lmbda was removed from DistillationConfig in this PR, so the xfailed num_items test was raising TypeError at construction — swallowed by the xfail marker, so it never exercised the num_items path. The trainer is always on-policy now, so the arg is unnecessary.
# Conflicts: # trl/experimental/distillation/distillation_trainer.py
cmpatino
approved these changes
Jul 21, 2026
kashif
pushed a commit
to kashif/trl
that referenced
this pull request
Jul 23, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Plan step 3.1 — stacked on PR #6460. First PR of Phase 3 (dataset contract → prompt-only). Part of #6449.
The trainer currently only accepts a
messages(conversational language-modeling) dataset. This PR additively teaches the collator the forward-looking prompt-only format — apromptcolumn of conversational messages, the same column GRPO consumes — as the first step toward making prompt-only the sole contract (later in Phase 3).Trainer
_DistillationCollatornow branches on the example's columns: apromptcolumn is treated as a completion-less prompt (the student generates on-policy), rendered and tokenized through the exact same path as the existingmessagesprompt (apply_chat_template(..., add_generation_prompt=True), prompt-budget truncation, every label masked). Amessagescolumn keeps its current prompt/last-assistant-turn split. Nothing about themessagespath changes.promptis registered in_set_signature_columns_if_needed(moving toward GRPO's pinned["prompt", "image", "images"], which lands at plan step 3.6).Scope
Conversational
prompt(a list of messages) only — the direct analog of the existing conversationalmessagessupport. Standard (plain-string) prompts arrive with the GRPO_tokenize_promptsadoption in Phase 5.Tests
test_collator_accepts_prompt_only_dataset: aconversational_prompt_onlybatch collates to a fully-masked, prompt-only tensor set (labels == -100,input_ids == prompts).test_train_runs_with_prompt_only_dataset: end-to-end prompt-only training, added as a strict xfail. Prompt-only datasets carry no dataset-completion tokens, sonum_items_in_batch(counted from the raw dataloader labels before on-policy generation replaces them) is0, making the losssum / 0 = NaN. This is the same denominator bug already documented at plan step 0.3; both un-xfail together when it is fixed at plan step 5.6.Verification (pytest)
Full
test_distillation_trainer.py: 38 passed, 2 xfailed (the pre-existingnum_items_in_batchxfail + the new prompt-only one). Collator test passes; the training xfail deterministically fails on the NaN loss.Risk
Low. Purely additive — the
messagespath is byte-for-byte unchanged, and the newpromptbranch reuses the same render/tokenize code. The only new runtime surface is column detection at the top of the collator loop.Note
Low Risk
Additive column branching in the collator; the
messagescode path is untouched and prompt-only rows reuse the same tokenization path.Overview
DistillationTrainer can now ingest GRPO-style prompt-only datasets: examples with a
promptcolumn (conversational messages, no assistant completion in the data) are collated like prompt-onlymessagesrows—all labels masked, completions left for on-policy generation. The existingmessagespath (prompt = everything before the last assistant turn) is unchanged._DistillationCollatordetectspromptvsmessagesat the start of each example; docs describe both contracts.promptis added to_set_signature_columns_if_neededso the column is not stripped by the trainer.Tests add
test_train_runs_with_prompt_only_datasetas a strict xfail: with no dataset completion tokens,num_items_in_batchstays 0 before generation and the loss becomes NaN—the same denominator issue as the existing on-policynum_items_in_batchtest, slated for plan step 5.6.Reviewed by Cursor Bugbot for commit 71883b9. Bugbot is set up for automated code reviews on this repo. Configure here.