[DistillationTrainer refactor] Emit prompt_ids / prompt_mask / completion_ids#6487
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.
…pport # Conflicts: # tests/experimental/test_distillation_trainer.py
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.
Item 20 (Group B) — stacked on PR 19. Part of #6449.
Third step of the layout migration (add → switch → delete). Emit GRPO's keys alongside the existing ones; nothing consumes them yet (item 21 switches consumers, item 22 deletes the old helpers).
prompt_ids(=prompts),prompt_mask(=prompt_attention_mask), andcompletion_ids(the completion region ofinput_ids; empty for the prompt-only collator).New test asserts
cat(prompt_ids, completion_ids) == input_idsand that the prompt keys mirror the existing tensors.Verified: 53 tests pass (distillation + server); ruff clean.
Note
Low Risk
Additive tensor keys only; loss and training logic still use the legacy fields, so behavior is unchanged aside from slightly larger batch dicts.
Overview
DistillationTrainer batch dicts gain GRPO-style
prompt_ids,prompt_mask, andcompletion_idsin addition to the existingprompts/prompt_attention_masklayout. Nothing reads the new keys yet; this is the “add alongside” step before consumers switch and old helpers are removed.The
_DistillationCollatorreturn value adds the three keys (prompt_ids/prompt_maskmirror the prompt tensors;completion_idsis the tail slice ofinput_ids, empty for prompt-only collate). After on-policy generation,_generate_with_modeland_store_completions_in_bufferpopulate the same fields on buffered microbatches.A new integration test captures inputs in
compute_lossand asserts the keys exist, thatcat(prompt_ids, completion_ids) == input_ids, and that the prompt keys matchprompts/prompt_attention_mask.Reviewed by Cursor Bugbot for commit 69bbafc. Bugbot is set up for automated code reviews on this repo. Configure here.