docs/engineering_insights.md

Engineering Insights

This project evolved through several failed and successful training designs. The useful lessons are summarized here as public engineering notes.

For expanded operational detail, see Training Playbook, Pipeline Hardening, and Evaluation Methodology.

1. Sparse Offline KD Hit A Ceiling

The earliest distillation path cached only a small top-k slice of teacher logits. That made training cheaper, but it discarded most of the teacher distribution. With a vocabulary of roughly 151K tokens and $k = 8$, the visible support was:

$$\frac{k}{\mathrm{\mid }V\mathrm{\mid }}=\frac{8}{\mathrm{151,665}}\approx 5.3\times {10}^{-5}=0.0053\mathrm{\%}\backslash frac\{k\}\{|V|\}\; =\; \backslash frac\{8\}\{151\{,\}665\}\; \backslash approx\; 5.3\; \backslash times\; 10^\{-5\}\; =\; 0.0053\backslash \%$$

The result was clear: top-k KD could perturb the student, but it did not transfer enough "dark knowledge" to reliably improve reasoning. Different alphas, epochs, and student initializations could not escape this sparse-signal ceiling.

The final fix was to use online KD: load teacher and student together, run both forward passes, and compute KL against the teacher's full vocabulary distribution.

2. Base Student Was Better Than Fighting An Aligned Space

Distilling into an already instruction-tuned student can cause destructive interference. The student's weights already encode one aligned behavior manifold, while the teacher's soft logits pull toward another. Training can look numerically stable while reasoning metrics regress.

The final path uses Qwen/Qwen3-1.7B-Base as the student. The base model has more plasticity, while the CE term and later SFT stage teach assistant formatting.

3. KD And Alignment Are Different Problems

Standardized benchmarks showed that KD can improve reasoning and calibration, but open-ended chat quality still needs alignment data.

The important diagnosis:

• A distillation failure means the student did not absorb the teacher's useful probability structure.
• An alignment gap means the student has capability, but the generation path is not yet trained to behave like a polished assistant.

The project therefore separates the pipeline into KD first, then SFT.

4. Assistant-Only Loss Masking Matters

A key bug class was assigning loss to chat formatting tokens instead of only assistant response content. If the model is trained to optimize structural tokens too heavily, it can learn formatting before substance.

The current tokenization path derives an assistant-only loss_mask, so:

• User prompts are context, not targets.
• Chat headers and separators are masked.
• Assistant response tokens are the only supervised targets.

This keeps training focused on semantic outputs rather than wrapper reproduction.

5. Sequence Packing Was The Main Throughput Win

The dataset contains many sequences shorter than the maximum context length. Dynamic padding wastes a large fraction of compute. First-fit decreasing sequence packing converted that waste into useful tokens.

Observed engineering outcome:

• Unpacked B200 online KD ran around the low-20K tokens/sec range in earlier probes.
• Packed B200 online KD reached roughly the mid-40K tokens/sec range after warmup.
• Packed utilization was close to full 4096-token bins.

The final code keeps packing deterministic and stores packing metadata in checkpoints so packed/unpacked resume mismatches fail loudly.

6. Full-Vocab KD Needed Token Chunking

Online KD preserves the full teacher distribution, but a full KL workspace at Qwen vocabulary scale is too large to materialize casually:

$$\text{KL workspace}\sim {\mathbb{R}}^{B\times S\times \mathrm{\mid }V\mathrm{\mid }}\backslash text\{KL\; workspace\}\; \backslash sim\; \backslash mathbb\{R\}^\{B\; \backslash times\; S\; \backslash times\; |V|\}$$

The solution is token-dimension chunking. The current implementation uses:

$$C_{\text{KD}}=2048C\_\{\backslash text\{KD\}\}\; =\; 2048$$

Larger chunks reduce loop overhead, but increase temporary memory pressure. The selected value is a practical B200-oriented balance for the 8B -> 1.7B workload.

7. Shape Churn And Synchronization Can Quietly Drain Throughput

Several performance bugs were not correctness bugs:

• Dynamic sequence lengths caused allocator churn.
• Repeated .item() calls forced CPU-GPU synchronization.
• Single-GPU DeepSpeed could add overhead when the model already fit comfortably.
• torch.compile added memory overhead, dynamic-shape graph breaks, recompile overhead, and checkpoint portability risk.

The final training loop favors stable shapes, fewer scalar syncs, fused AdamW when available, FlashAttention when available, and Liger kernels where they do not conflict with KD logits.

8. Evaluation Requires Controlled Comparisons

Raw completion and chat-template evaluation activate different behavior. A base model can perform well in raw mode and poorly under chat markup. A chat-aligned model can underperform on raw continuation-style tasks if the benchmark asks for direct option likelihoods.

The project uses both controls:

• Raw-to-raw comparisons isolate distilled base capability.
• Chat-to-chat comparisons estimate template robustness and assistant-format alignment.

This distinction avoids blaming KD for failures that belong to alignment or benchmark formatting.

9. Post-KD SFT Is Not Optional For Assistant Quality

KD transfers probability structure; it does not guarantee careful behavior, refusal policy, calibrated uncertainty, or code reliability. Targeted SFT was added to address:

• Confident hallucination in open-ended answers.
• Persona and identity consistency.
• Repetition loops.
• Chat-format stability.
• Practical assistant presentation.

Preference training or DPO would be the natural next layer if the project continues beyond the current release.

10. Training Loss Is Not The Release Gate

Several development runs looked numerically healthy while downstream benchmarks moved in the wrong direction. That pattern is expected when the training objective is only a proxy for the release objective.

Useful release gates:

• Standardized benchmarks.
• Raw and chat controls.
• Mismatch inspection.
• Qualitative prompts after benchmark checks.
• Weight and checkpoint structure audits.

Held-out KD validation loss is important, but it cannot prove that the model improved on math, code, multiple-choice reasoning, or assistant behavior.

11. Fail-Fast Beats Silent Recovery

The pipeline hardened around a simple rule: corrupt artifacts should stop the run.

Examples:

• Missing teacher-logit shards fail instead of becoming zero tensors.
• Tokenization with zero usable rows fails immediately.
• Shard schema mismatches are rejected.
• Packed/unpacked checkpoint resume mismatches are rejected.
• Stale evaluation outputs are cleaned before new scores are written.

This makes errors louder, but it keeps published numbers trustworthy.

12. Public Docs Should Preserve Decisions

A release-quality project should expose durable engineering conclusions:

• why online KD replaced offline top-k KD,
• why assistant-only masking matters,
• why raw/chat evaluation controls are required,
• why sequence packing changed throughput,
• why SFT remains necessary after KD,
• why checkpoint and provenance checks exist.

That level of detail is enough for technical readers without turning the documentation into a chronological run journal.