Self-distillation + A100 campaign prep - #1
Merged
Conversation
Train the drafter on the target's OWN greedy generations so the next-token label equals the target's argmax — exactly what decode-time acceptance measures — directly optimizing acceptance length (the speedup gate). - pe.distill: generate target greedy response per prompt, featurize [prompt | response], shard in FeatureDataset layout + per-example prompt_len - FeatureDataset(with_prompt_len): yields prompt_len; flags self_distilled caches - mtp_backward(prompt_len): masks prompt-region labels out of the loss (only response labels equal the target argmax) — exact valid-slot count - train: auto-detects self-distilled caches and masks the prompt region - tests: response labels == target argmax; masking counts; backward accepts it - make distill target
Batched greedy generation (HF .generate, left-padded) makes A100 data-gen ~10x faster and avoids long single-sequence spot exposure. Left-padding can flip occasional near-tie tokens vs unpadded decoding, so supervision is taken from the unpadded teacher-forced argmax (same forward as the features) rather than the generated token — labels are exact regardless of generation method. - distill.distill_batch: left-pad a prompt batch, batched .generate, per-example unpadded featurization (positions stay inference-exact) - _featurize returns per-position labels = target argmax (features+labels from one consistent forward); also tightens the single-sequence path - mtp_backward(teacher_labels): supervise on teacher_labels[tgt-1] when present, else next token (human-text caches unchanged) - FeatureDataset yields labels; train threads them through - --batch-size CLI flag; single-sequence remains the default - tests: batched/single labels == unpadded argmax (real-model verified 0 label mismatches); teacher-label path == next-token path when self-consistent 31->32 tests, ruff clean.
run_bench timed the prefix-recompute loop (generate_speculative / vanilla_generate); the headline now uses the KV-cached path we actually ship — generate_speculative_cached vs vanilla_generate_cached — so reported speedup/calls-per-token reflect production. --recompute keeps the old loop for the lossless-equivalence reference. Verified lossless on the toy model for chain and tree.
The cached one-forward loop emitted on_commit one token at a time, so the streaming demo could never show multi-token bursts on the path we ship — the 'see the difference' artifact only worked on the slow recompute loop. Emit each step's accepted burst in a single on_commit call (display only; output_ids unchanged, still lossless — verified on toy at acceptance 5.7 -> bursts of 6). - serve.generate_speculative_cached: batch on_commit per step - scripts/record_demo.sh: asciinema -> agg GIF of naive vs tree - make demo (live) / make record-demo (GIF), TARGET/CKPT/PROMPT overridable
scripts/azure_a100_setup.sh: provision a single A100 80GB spot VM with a budget guard (alert at $80 of the $100 credits), NVIDIA driver extension, and the on-VM bring-up steps; deallocate to stop hourly billing. scripts/run_campaign.sh: the end-to-end A100 run — batched self-distill -> prompt-masked train -> cached-loop benchmark -> side-by-side GIF, all parameterized (TARGET/EXAMPLES/BATCH/LAYERS/DEPTH).
Mid-epoch checkpointing for cloud GPUs where an epoch can outlast the mean time-between-preemptions. Training saves the model + training position every save_every updates (atomic write + optional on_save volume-commit hook) and a restarted process fast-forwards into the in-progress epoch via the deterministic dataset iterator; the optimizer is re-initialised (Adam's first post-reset step is self-normalised, far cheaper than persisting ~5GB of moments). A config signature guards against resuming a mismatched checkpoint. Three new tests cover resume-continues, finished-run no-op, and mismatched-config-starts-fresh.
infra/modal_campaign.py runs the distill -> train -> bench pipeline on a Modal A100 (env-overridable target/GPU), with a detach-safe single-container campaign, a resume-from-cache 'finish' entrypoint (mid-epoch checkpoint + retries for preemption resilience), and a read-only 'diagnose' entrypoint. bench/diagnose_acceptance.py measures the drafter's per-depth top-1 accuracy against the target's own greedy argmax, decomposing acceptance length into the strong depth-0 prediction vs the data-hungry deeper draft positions.
Scaling the pipeline to an fp16 Qwen2.5-14B target on an A100, the trained 4-layer drafter crosses break-even: at 15 epochs the parallel-tree loop runs 1.045x vanilla (target_calls/token 0.82 < 1.0), with acceptance climbing monotonically with training (1.345 -> 1.447 @6ep -> 1.677 @15ep). README adds the 14B results table, a per-depth diagnostic explaining the trajectory, and a note framing the remaining gap to a large speedup as training scale, not architecture. Ignore distill_cache*/.
Adds bench/demo_race.py (record both decoders once, replay each stream, composite side by side) and bench/sidebyside.py (stitch two terminal GIFs time-aligned). The Modal gif entrypoint probes several prompts on the saved 14B drafter, keeps the best-speedup one, and renders the side-by-side demo. README embeds docs/demo.gif and the 14B benchmark plot, and documents that speedup is much larger on structured/code prompts (~1.3x, acceptance ~1.9-2.0) than the mixed-prompt average. Also fixes the streaming demo to decode incrementally so inter-token spaces survive.
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
1 participant
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.
Why
The system is lossless and complete, but the headline — a clearly-visible wall-clock speedup — was gated by drafter acceptance (~1.4, need ≳2.0). Acceptance measures "did the drafter predict the target's argmax?", yet we trained on human-written text (label = human token ≠ target argmax). This PR closes that gap with self-distillation and prepares the A100 campaign to realize the speedup on a big fp16 target (where speedup ≈ acceptance, since the tiny drafter's cost is negligible next to each target forward).
What
pe.distill): generate the target's own greedy response per prompt, cache features over[prompt | response]. Label = target argmax by construction.mtp_backward(prompt_len=…)): only response-region labels (which equal the target argmax) carry the acceptance signal.--batch-size): ~10× faster data-gen, keeps A100 spot exposure short; single-sequence remains the default.run_bench.pynow times the cached one-forward loop (generate_speculative_cachedvsvanilla_generate_cached);--recomputekeeps the old reference.on_commit, so streaming shows multi-token bursts (verified bursts of 6 at acceptance 5.7, still lossless).make demo(live) +scripts/record_demo.sh(GIF).scripts/azure_a100_setup.sh(A100 spot VM + budget guard) andscripts/run_campaign.sh(distill → train → bench → GIF).Tests
32 tests (was 27), ruff clean. New tests prove: response/batched labels == unpadded target argmax; prompt-masking slot counts; teacher-label path == next-token path when self-consistent.
Not in this PR
The A100 run itself (Qwen2.5-14B fp16) — gated on a free local validation that self-distillation lifts acceptance above the 1.39 baseline.
🤖 Generated with Claude Code