# 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}} $$ For the final Qwen3 run: $$ \alpha = 0.3,\quad T = 2.0,\quad C_{\text{KD}} = 2048,\quad S_{\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}{|V|} = \frac{8}{151{,}665} \approx 5.3 \times 10^{-5} = 0.0053\% $$ 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 |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 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}} = 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,\quad A=2,\quad B_{\text{eff}}=8,\quad L_{\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: ```text 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.