[GKD] Buffer Implementation for Distillation Trainer

[cmpatino]

Implement Buffer for Distillation Trainer (GOLDTrainer)

Implement generation buffering and multi-generation support for GOLDTrainer

Add a prompt-level generation buffer that decouples generation from the
optimization steps. We adopt a buffer similar to GRPO to generate all rollouts for all mini-batches within an optimization step, leveraging parallel inference engines. This means each worker handles a buffer of per_device_train_batch_size * gradient_accumulation_steps.

Buffer Details

We allow multiple rollouts per prompt, following Thinking Machine’s Tinker example. The number of rollouts per prompt is determined by the num_generations parameter. To keep the effective batch size constant, we introduce the generation_batch_size parameter, which controls how many unique prompts we pass to the inference engine. We enforce generation_batch_size = per_device_train_batch_size * gradient_accumulation_steps // num_generations to ensure the effective batch size is invariant across setups.

Benchmarks

We can replicate Thinking Machine’s results using both non-Liger and Liger losses, achieving a 3x speedup on a setup with 8 training nodes in colocate mode.

Phase	Tinker (s)	TRL (s)
Sampling	329.83	130
Loss	37.96	-
Training	98.69	38
Total	492.28	173

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

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.

---

Note

Medium Risk
Changes the GOLD training loop and dataloader sampling semantics to buffer generations across optimizer windows and support multiple rollouts per prompt, which can affect training dynamics and correctness across distributed setups.

Overview
Adds an optimizer-window generation buffer to GOLDTrainer, decoupling on-policy rollout generation from individual gradient-accumulation microsteps. Training now samples a full window batch via a RepeatSampler, generates (optionally via vLLM) once per window, and reuses buffered slices across accumulation steps.

Extends GOLDConfig with num_generations and generation_batch_size (with strict partitioning validation) and adds teacher_model_revision; updates model/teacher revision handling in gold.py and teacher model instantiation. Generation/label masking is refactored to correctly handle left-padded prompts and completion-only outputs, and Liger+ZeRO-3 gets an explicit parameter gather context.

Updates docs to describe buffering behavior and config changes, and adds/adjusts tests to cover left-padding retokenization and prompt-length masking expectations.

^{Written by Cursor Bugbot for commit abfadc1. This will update automatically on new commits. Configure here.}

[HuggingFaceDocBuilderDev]

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.

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[qgallouedec]

I haven't reviewed it in detail; I have a general idea of what it's about, but I'm leaving the implementation mostly up to you. In future PRs, we can try to align it better with the rest of the codebase, but what matters most right now are the results you're getting.

Make sure to run make precommit to make to CI happy

[qgallouedec]

just ensure the CI is green before merging