[gold] Implement seq_kd in GOLDTrainer - #5725
Conversation
|
Thanks @roycho96, and sorry for the long silence; this slipped through our review backlog The change is in good shape: The catch: #5969 (VLM support) has since reworked |
GOLDConfig documents seq_kd but the flag was never read. Off-policy slices always reused the dataset completion. Now, when seq_kd=True, _fill_buffer routes off-policy slices to _generate_seq_kd_for_slices, which batches all prompts into a single teacher.generate call and buffers the completions for training, matching GKDTrainer semantics. Two completion mappings: with ULD and a separate teacher tokenizer the teacher output round-trips through text into the student vocab. With a shared tokenizer the raw teacher ids pass through, which preserves EOS and avoids drift at BPE boundaries. Reconciled with the VLM support added in huggingface#5969. The VLM path buffers raw examples lazily and has no prompt tensors when _fill_buffer runs, while seq_kd feeds prompt ids to teacher_model.generate, so combining the two would need teacher-side image preprocessing and a synthetic example rebuild. GKDTrainer has no VLM support, so parity does not require it. seq_kd=True with a vision dataset now raises a ValueError in __init__ next to the existing VLM guards. A VLM student trained on a text-only dataset still works with seq_kd. Also fix a latent dtype bug in _process_completions_to_buffer: an empty completion produced torch.tensor([]) as float32, corrupting the dtype of input_ids and labels. Force dtype=torch.long.
507d819 to
691052e
Compare
![cursor[bot]](https://cdn.github.com/in/1210556?s=60&v=4){kind=small}
Thanks for the detailed review! Rebased onto main to pick up #5969, with the seq_kd routing moved into the text branch of _fill_buffer. |
When pad_token_id equals eos_token_id, stripping trailing pad tokens from the teacher completion also removes the terminating EOS, so the student never trains on a stop token. Re-add EOS when tokens were stripped and pad equals eos.
Decoding teacher completions with skip_special_tokens=True removes the teacher EOS before re-encoding into the student vocab, so terminated completions never carried a stop token and a completion consisting only of EOS became empty. Append the student EOS whenever the raw teacher completion contains the teacher EOS.
In the cross-tokenizer seq_kd path the teacher prompt was decoded with skip_special_tokens=True, so the teacher generated from bare concatenated text without the student chat markers. That is out of distribution for an instruct teacher, and the ULD loss then scored the resulting completions under the marker-full text (original_prompt_text, decoded with skip_special_tokens=False), so generation and loss disagreed on the teacher conditioning. Feed the teacher tokenizer the marker-full prompts_text already computed in _generate_seq_kd_for_slices. Generation and loss now share the same id source, the same skip_special_tokens=False decode, and the same add_special_tokens=True, matching the on-policy path as well. Add a test that locks the teacher tokenizer input to the skip_special_tokens=False decode of the student prompt ids.
When the teacher tokenizer has no eos_token_id, the cross-tokenizer path substituted the student's EOS id, an arbitrary token in the teacher vocab, so teacher generation could stop mid-sentence on a random token and the same bogus id leaked into pad_token_id. Drop the fallback: with eos_token_id None, generate() runs to max_new_tokens, which is the correct no-EOS semantics, and __init__ already assumes KD teachers define an EOS. With teacher_eos now possibly None, the guard that restores the student EOS after decode must check it explicitly: row == None evaluates to the Python bool False and .any() raises AttributeError. Skipping the append there is also semantically right, since without a teacher EOS no completion terminated on one and there is no stop signal to restore.
The non-vLLM on-policy path called _build_sequence_batch without an attention mask, taking the helper's identity-based pad masking branch. When pad_token_id == eos_token_id (SmolLM, Qwen), that set the terminating EOS label to -100 on every generated row, so the student never received stop supervision, and it zeroed attention on real pad-id tokens inside the prompt, such as chat turn delimiters. Build the attention mask with the first-EOS positional masking already used by GRPO and RLOO, identical to the GKD fix, and pass it to _build_sequence_batch explicitly, routing this caller through the same helper branch the buffered path already uses. The terminating EOS stays attended and labeled, only padding strictly after it is masked, prompt columns take the prompt attention mask. The buffered seq_kd path is untouched.
This reverts commit 1840ea4. The on-policy EOS masking fix patches pre-existing generate_on_policy_outputs, not the seq_kd feature this PR adds, so it moves to its own fix PR against main. This branch keeps only the seq_kd-specific changes.
The teacher previously generated from the full untruncated prompt and relied on the post-hoc truncation in _process_completions_to_buffer, which keeps the start of the prompt by default and drops the generation marker at the end. Budgeting the prompt keep-end before generation makes the teacher generate from and train on the same context, keeps prompt plus completion within max_length, and mirrors the on-policy budgeting idiom.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit f2b2715. Configure here.
|
I split a separate PR to fix gkd eos masking when pad equals eos |
|
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. |
4 participants
What does this PR do?
GOLDConfig.seq_kdand a docstring saying it was "inherited fromGKDConfig" have existed since #4349, but the trainer never actually checked the flag. This PR implements it.Behavior
_fill_bufferpartitions slices into three exclusive groups: on-policy, seq_kd, dataset. GKD runs the two branches in sequence, soseq_kd=Truewithlmbda > 0has the teacher generate output that the student immediately overwrites. GOLD picks one path per slice._generate_seq_kd_for_slicescollects all seq_kd prompts in a buffer fill and runs a single batchedteacher.generate().Also fixed
When the teacher emits only EOS, the cross-tokenizer round-trip gives an empty completion (
[]). That triggers a latent bug in_process_completions_to_buffer.torch.tensor([])defaults to float, the latertorch.catthen promotes the whole row, andcompute_losscrashes at the embedding lookup.mainnever hits this because vLLM always returns at least one token.Fix: pass
dtype=torch.longtotorch.tensor.Regression test:
test_process_completions_to_buffer_handles_empty_completion.Tests
In
tests/experimental/test_gold_trainer.py. Six unit tests cover same-tok content, cross-tok round-trip, single batched generate call, and slice routing with seq_kd on and off. One@pytest.mark.slowend-to-end smoke confirmstrainer.train()runs without crashing.Before submitting
AI writing disclosure
We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.
@kashif @qgallouedec @albertvillanova
Note
Medium Risk
Changes core training buffer routing and teacher generation for off-policy steps; well covered by tests but affects distillation behavior when seq_kd or lmbda mix on-policy and off-policy data.
Overview
Implements
GOLDConfig.seq_kd, which existed on the config but was never honored in the trainer. Off-policy microbatches can now use teacher-generated completions (sequence-level KD) instead of dataset labels._fill_bufferassigns each gradient-accumulation slice to one of three paths: on-policy student generation, seq_kd teacher generation, or the original dataset batch (VLM lazy collation unchanged)._generate_seq_kd_for_slicesbatches all seq_kd prompts into a singleteacher.generate(), budgets prompts with keep-end truncation like on-policy, and maps completions into the student buffer—raw teacher token ids when tokenizers match, or decode/re-encode for ULD cross-tokenizer (with EOS/pad edge cases handled).seq_kd=Truewith vision datasets is rejected at init.Also fixes
_process_completions_to_bufferso empty completions usedtype=torch.long(avoids float promotion and embedding crashes on EOS-only teacher output).Tests cover routing, batched teacher calls, cross-tokenizer behavior, empty completions, vision+seq_kd rejection, and a slow end-to-end
trainer.train()smoke withseq_kd=True.Reviewed by Cursor Bugbot for commit a30830f. Bugbot is set up for automated code reviews on this repo. Configure here.