docs/training_playbook.md

Training Playbook

This page captures the practical training lessons behind Quintus. It focuses on the engineering decisions that made the final online-KD run stable, reproducible, and fast enough to complete on large single-GPU hardware.

Core Objective

The training objective combines assistant-token cross entropy with teacher-student KL divergence:

$${\mathcal{L}}_{\text{total}}=\alpha {\mathcal{L}}_{\text{CE}}+(1-\alpha ){\mathcal{L}}_{\text{KD}}\backslash mathcal\{L\}\_\{\backslash text\{total\}\}\; =\; \backslash alpha\; \backslash mathcal\{L\}\_\{\backslash text\{CE\}\}\; +\; (1\; -\; \backslash alpha)\backslash mathcal\{L\}\_\{\backslash text\{KD\}\}$$

For the final Qwen3 run:

$$\alpha =0.3,T=2.0,C_{\text{KD}}=2048,S_{\max }=4096\backslash alpha\; =\; 0.3,\backslash quad\; T\; =\; 2.0,\backslash quad\; C\_\{\backslash text\{KD\}\}\; =\; 2048,\backslash quad\; S\_\{\backslash max\}\; =\; 4096$$

In this codebase, $\alpha$ is the cross-entropy weight. Lower $\alpha$ gives the teacher distribution more influence. Higher $\alpha$ gives hard assistant targets more influence.

Why Online KD Replaced Offline Top-K KD

The early pipeline precomputed only a small top-k slice of the teacher distribution. That made storage and training cheaper, but it created a hard information ceiling.

With a Qwen vocabulary around 151K tokens:

$$\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 \%$$

That sparse signal was enough to disturb student weights, but not enough to reliably transfer deeper reasoning behavior. Several development probes changed alpha, epochs, and student initialization; the same ceiling remained.

The final online path removes that bottleneck. Teacher and student run together, and the KL term is computed from the live full-vocabulary teacher distribution.

Memory Shape To Respect

Full-vocabulary KD is dominated by logits:

$$\text{student\_logits},\text{ teacher\_logits}\in {\mathbb{R}}^{B\times S\times \mathrm{\mid }V\mathrm{\mid }}\backslash text\{student\backslash \_logits\},\backslash \; \backslash text\{teacher\backslash \_logits\}\; \backslash in\; \backslash mathbb\{R\}^\{B\; \backslash times\; S\; \backslash times\; |V|\}$$

At Qwen vocabulary scale, increasing micro-batch size by one can add many GiB of temporary memory pressure. Effective batch size is not the same as memory cost. Peak memory is mostly driven by micro-batch size, sequence length, vocabulary width, activation storage, and the backward pass.

Useful rule:

$$B_{\text{eff}}=B_{\mu }\times AB\_\{\backslash text\{eff\}\}\; =\; B\_\{\backslash mu\}\; \backslash times\; A$$

Keeping $B_{\mu}$ lower and $A$ higher is often safer than a large micro-batch with the same effective batch size.

Token Chunking

A naive full-vocabulary KL implementation materializes too much temporary state. Quintus computes KD over token chunks:

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

Larger chunks reduce loop overhead but increase temporary memory. Smaller chunks save memory but can add kernel-launch and Python overhead. The final value is a B200-oriented balance for the 8B -> 1.7B workload.

Sequence Packing

Sequence packing was the largest throughput win in development probes.

The packing strategy:

• Sort samples by length descending.
• Pack samples with deterministic first-fit decreasing binning.
• Insert EOS separators between samples.
• Set separator loss_mask = 0.
• Optionally mask the first token after each separator.
• Build attention_mask from true packed length, not from token identity.

The attention-mask detail matters because Qwen tokenizers can share EOS-like IDs with padding behavior. Deriving attention from input_ids != pad_token_id can accidentally mask real EOS separators inside packed rows.

Packing probes showed an unpacked B200 online-KD baseline around the low-20K tokens/sec range. Packed training reached roughly the mid-40K tokens/sec range after warmup. The final Qwen3 profile uses the same design principle with a conservative 8B -> 1.7B batch shape.

B200-Oriented Final Shape

The Qwen3 config is intentionally conservative:

$$B_{\mu }=4,A=2,B_{\text{eff}}=8,L_{\text{pack}}=4096B\_\{\backslash mu\}=4,\backslash quad\; A=2,\backslash quad\; B\_\{\backslash text\{eff\}\}=8,\backslash quad\; L\_\{\backslash text\{pack\}\}=4096$$

Runtime choices:

• gradient_checkpointing = false
• compile_model = false
• fused_adamw = true
• sequence_packing.enabled = true
• FlashAttention-2 when available
• Liger kernels for compatible Qwen-family operators

The main reason is the 8B teacher plus 1.7B student online-KD footprint. A smaller teacher/student pair can use larger micro-batches, but the release workload reserves more headroom.

Kernel Choices

FlashAttention-2 is the preferred stable attention path when available.

Liger kernels are useful for Qwen-family training, but KD places an important constraint on fusion:

• Safe to fuse: RMSNorm, RoPE, SwiGLU.
• Avoid for KD: fused linear cross entropy that hides raw student logits.

The KD loss needs raw student logits to compute teacher-student KL. Any optimization that bypasses logits entirely can break the objective.

Why torch.compile Stayed Off

torch.compile can be useful for some SFT paths, but it was not the production choice for final KD.

Observed risks:

• Large Inductor memory overhead.
• Warmup cost on short-lived cloud instances.
• Dynamic-shape graph breaks from variable sequence lengths.
• Recompile overhead that reduced cumulative throughput in probes.
• _orig_mod. prefixes in saved checkpoints if compiled modules are not unwrapped before saving.
• Limited benefit after FlashAttention and Liger already fuse the major kernels.

For this workload, stable eager execution with targeted kernels was more predictable than compiler-driven fusion.

DataLoader And Cloud Stability

Large worker counts can improve throughput on local systems, but notebook and cloud environments can deadlock through multiprocessing queues, IPC limits, or shared-memory pressure.

Practical policy:

• Start with conservative worker and prefetch settings.
• Treat a silent training hang as a DataLoader candidate, even when GPU utilization remains high.
• For some cloud notebook runs, dataloader_workers = 0 was the most stable choice.
• For the release config, dataloader_workers = 8 and prefetch_factor = 2 are a controlled default, not a universal rule.

Checkpointing And Resume

Cloud GPUs are preemptible and notebook sessions disappear. The training loop therefore treats checkpointing as a core training feature, not an afterthought.

Important design points:

• best is selected from validation loss where available.
• last is saved for final-state inspection.
• Step checkpoints can resume mid-epoch.
• Scheduler state is saved.
• Optimizer state may be intentionally omitted for very large runs to avoid massive checkpoint overhead.
• Resume semantics distinguish initialization from a completed checkpoint and continuation from an interrupted checkpoint.

This avoids the common trap where resume_from_checkpoint silently starts from the wrong phase or stale state.

Provenance Rules

The pipeline is strict about artifact compatibility:

• Tokenizer vocabulary sizes must match the model contract.
• Teacher-logit metadata must match expected temperature, sample count, max sequence length, and tokenizer/model identity.
• Dataset fingerprints are preferred over path equality because paths are machine-local.
• Tokenizer fingerprints can drift across library versions, so hard checks should focus on vocab-size and schema invariants.

The principle is simple: train only when artifacts prove they belong together.

Dataset Sampling

Taking the first N valid streamed examples can bias a run if the upstream dataset is ordered by source, task, difficulty, or language. Later configs added stream shuffling before selection.

The config uses a non-default seed:

stream_shuffle_seed = 25
split_seed = 25

The number is intentionally explicit. Reproducibility needs stable seeds; it does not require the overused value 42.

Practical Watchpoints

During a run, these signals matter more than a single loss number:

• Loss stays finite from the first logging window.
• CE and KD move in plausible ranges.
• Rolling throughput remains stable after warmup.
• GPU memory is high but not near an unpredictable OOM edge.
• Validation loss is computed on the intended holdout.
• Saved checkpoints load in standard Transformers and vLLM paths.
• Downstream benchmark results agree with the training story.

Held-out KD loss is useful, but it is not the release gate. Standardized benchmarks and qualitative checks must decide whether the checkpoint improved the target behavior.