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}}
$$

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.