Two separate artifacts come out of this stage:
teacher_outputs.jsonl— the teacher's summary text. Feeds sequence-level KD (CE loss).logits/*.npz— top-K teacher logits per token position. Feeds logit-level KD (KL loss).
You can do (1) alone. You cannot do (2) alone (you need a target sequence to score).
| Strategy | Setting | Effect on the student |
|---|---|---|
| Greedy | temperature=0 |
Sharpest, most consistent targets. Student learns the teacher's mode. Lowest diversity; can amplify teacher tics. |
| Low-temp sampling | T=0.7, top_p=0.9 |
Small diversity injection. Good default. |
| High-temp sampling | T=1.0+ |
Diverse but noisier; more teacher hallucinations enter your training set |
| Best-of-n + reranker | n=4, pick by a quality scorer | Highest quality data, 4x cost. This is "rejection sampling / RFT" and it materially beats plain sampling |
| Multiple samples kept | n=4, keep all 4 | Data augmentation: same doc, 4 targets. Teaches the student the distribution of good summaries, not one point |
Recommendation for Saransh: n=4 samples at T=0.8, top_p=0.95, then keep the best 1–2 by a
cheap filter (length constraint + no-hallucination heuristic + teacher self-score). This is
rejection-sampling distillation and it's usually the single biggest quality lever in the whole
pipeline — bigger than the choice of KD loss.
Greedy targets have zero entropy — every document maps to exactly one summary. The student learns a deterministic map, which is fine for a product but makes the CE loss floor artificially low and hides underfitting. Sampled targets keep some entropy and act as label smoothing.
Generating 100k summaries with HF generate at batch 8 takes days. vLLM with continuous batching
and prefix caching takes hours.
python -m vllm.entrypoints.openai.api_server \
--model Qwen/Qwen2.5-7B-Instruct \
--dtype bfloat16 \
--max-model-len 8192 \
--gpu-memory-utilization 0.90 \
--enable-prefix-caching \
--max-num-seqs 256Then src/data/generate_teacher.py hits it with an async client at concurrency 64.
Rough numbers on 1× A100-80GB, 7B bf16, ~1,500-token inputs / 150-token outputs:
- HF
generate, batch 8: ~1.5 examples/s → 100k = 18 hours - vLLM, prefix caching on: ~25–40 examples/s → 100k = ~45–70 minutes
Prefix caching matters because your system prompt + instruction is identical across examples;
vLLM reuses that KV cache. Put the shared text first (see docs/03 §4).
Teacher inference is the most expensive step. It will be interrupted.
- Write results with
idas the key, append-only JSONL,flush()every batch. - On restart, load existing ids into a set and skip them.
- Never overwrite the output file.
--resumeis the default ingenerate_teacher.py. - Record
teacher_model,teacher_revision(the exact HF commit sha), and all sampling params in every record. Six months later you will need to know which checkpoint produced a given row.
This is the part that surprises people.
Full logits for one 1,024-token sequence in fp16:
1024 × 151,936 × 2 bytes = 311 MB
For 50,000 examples: 15.5 TB. Not feasible.
The teacher's distribution is extremely peaked. Empirically, for summarization, the top-64 tokens hold >99% of the probability mass at almost every position. So store, per position:
topk_ids:int32[K]— vocabulary indicestopk_logits:float16[K]— their logits- (optionally)
residual:float16— log-sum-exp of everything outside the top-K, so you can reconstruct a proper normalized distribution with a single "everything else" bucket.
Storage per token at K=64: 64×4 + 64×2 + 2 = 386 bytes.
50,000 examples × 200 target tokens × 386 B = 3.9 GB ← completely fine
50,000 examples × 1,024 tokens × 386 B = 19.8 GB ← fine on a data disk
Only store logits for positions you compute loss on — i.e. the assistant/summary tokens, not the input document. For summarization the target is ~10% of the sequence, which is a 10x saving. This is the single most important storage optimization.
| K | Mass captured (typical) | Notes |
|---|---|---|
| 8 | ~95% | Too aggressive; loses the tail structure that is the dark knowledge |
| 32 | ~98.5% | Acceptable minimum |
| 64 | ~99.3% | Recommended default |
| 128 | ~99.7% | Marginal gain, 2x storage |
| full | 100% | Only viable with online teacher |
Measure it, don't guess: capture_logits.py --report-mass prints the mean captured mass so you can
pick K for your data. If your teacher is unusually high-entropy (very open-ended prompts), you may
need K=128.
Once you truncate to top-K, Σ p_i < 1. You have three options:
- Renormalize over top-K (
p_i / Σ_{topK} p_j). Simple, standard. Slight bias: it pretends the tail doesn't exist, which makes the target sharper than the teacher actually is. - Renormalize over the union of teacher-top-K and student-top-K. More faithful — it includes
positions where the student is putting mass that the teacher doesn't. Requires the student's
top-K at loss time. This is what
losses.pydoes whenunion_topk: true. - Add a residual bucket — one extra "all other tokens" class holding
1 - Σ_{topK} p. The student's mass on all non-top-K tokens is summed into a matching bucket. Mathematically the cleanest truncation of KL, and cheap. Recommended.
Option (1) is a common source of the symptom "distillation helps less than expected": you are training the student to be more confident than the teacher, which increases hallucination.
| Offline (precompute) | Online (teacher in memory) | |
|---|---|---|
| GPU memory | student only | student + teacher |
| Disk | 4–20 GB | none |
| Speed per step | fast (just I/O) | ~2x slower (extra forward pass) |
| Flexibility | K and target sequences are frozen | can distill on any sequence, including student-generated |
| Needed for GKD/on-policy | no — impossible | yes |
Use offline for Stages 1–2, switch to online for Stage 3 (on-policy).
The logits you store must be reproducible by the same teacher at training time, or your loss is noise. Sources of drift:
- Batch size changes logits. Different padding → different kernel reduction order → last-bit
differences. Harmless in magnitude (1e-3 in logit space), but it means you cannot unit-test for
exact equality; test
allclose(atol=1e-2). - Padding side and attention mask. If a padded position leaks into attention, logits are
garbage. Always pass
attention_mask. Test: same example alone vs in a padded batch → probs should match to ~1e-3. - Flash attention vs SDPA vs eager give slightly different numerics. Pin
attn_implementationin the config and record it. - dtype. Capture in bf16 or fp16 consistently, and store as fp16. bf16 has 8 mantissa bits; the logit quantization error is ~0.5% — irrelevant next to a temperature of 2.
use_cache=Trueduring a scoring forward pass can change numerics vs a single full pass. Useuse_cache=Falsefor teacher scoring.
logits[:, t, :] predicts token t+1. When you align teacher and student:
# both must be shifted the SAME way
s_logits = s_out.logits[:, :-1, :] # predictions for positions 1..T-1
t_logits = t_logits_stored[:, :-1, :]
labels = input_ids[:, 1:] # the actual next tokensIf you shift one and not the other, the student learns to predict the current token — which is trivially easy, so your loss will look great and your model will be useless. Detection: if training loss drops below ~0.05 within a few hundred steps, suspect this immediately.
# (a) generate teacher summaries via vLLM server
python src/data/generate_teacher.py \
--corpus data/processed/corpus.jsonl \
--out data/teacher/teacher_outputs.jsonl \
--model Qwen/Qwen2.5-7B-Instruct \
--n-samples 4 --temperature 0.8 --top-p 0.95 --max-new-tokens 384 \
--concurrency 64 --resume
# (b) capture top-K logits for the kept targets
python src/data/capture_logits.py \
--pairs data/teacher/teacher_outputs.filtered.jsonl \
--teacher Qwen/Qwen2.5-7B-Instruct \
--out-dir data/teacher/logits \
--topk 64 --batch-size 8 --report-mass --targets-onlyNext: 05 — data quality