docs/architecture.md

# Architecture

Quintus is built as a two-stage model development pipeline:

1. Online full-vocabulary knowledge distillation from a larger Qwen3 teacher into a Qwen3-1.7B base student.
2. Targeted SFT to improve instruction-following behavior, persona consistency, and generation stability.

![Quintus Architecture](../assets/quintus_architecture.png)

## Core Training Path

The main training entry point is `src/train.py`. It supports three phases:

- `sft`: Cross-entropy training on assistant response tokens.
- `kd`: Offline top-k teacher-logit distillation, retained for compatibility and provenance checks.
- `online_kd`: The final preferred path. Teacher logits are produced live during the student forward pass.

The final KD objective is implemented in `src/losses.py`:

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

For the final run, $\alpha = 0.3$ and $T = 2.0$. In this codebase, $\alpha$ is the cross-entropy weight. The complementary weight is assigned to the KD term.

## Data Flow

`src/download.py` prepares the training data. It handles both pre-tokenized rows and raw instruction data. For raw rows, it normalizes common conversation schemas, applies the tokenizer chat template, and builds an assistant-only `loss_mask`.

Important details:

- Prompt and formatting tokens are masked out.
- Assistant response tokens receive loss.
- Samples longer than `max_seq_len` are rejected rather than silently truncated.
- The tokenizer contract is later validated to avoid teacher/student vocabulary mismatches.

## Sequence Packing

`src/sequence_packing.py` implements deterministic first-fit decreasing packing. It places multiple shorter samples into fixed-length bins, separated by EOS tokens.

Packing properties:

- Training split is packed; validation can remain unpacked for interpretability.
- Bins are fixed at `pack_length = 4096` in the final profile.
- EOS separators have `loss_mask = 0`.
- The first token after a separator is optionally masked to avoid cross-sample target leakage.
- Attention masks are built from the true packed length, not by comparing token IDs against `pad_token_id`.

The attention-mask detail is important because Qwen tokenizers can reuse EOS-like IDs in ways that make token-identity-derived padding masks unsafe.

## Online KD Memory Strategy

Full-vocabulary KD is expensive because both student and teacher produce logits shaped as:

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

The implementation keeps this feasible by chunking along the token dimension with:

$$
C_{\text{KD}} = 2048
$$

Each chunk computes the teacher softmax, student log-softmax, and masked KL contribution, then accumulates the result. This preserves the dense teacher distribution while avoiding a single large KL workspace.

## Validation, Provenance, And Safety Checks

Several modules exist to prevent silent training corruption:

- `src/provenance.py`: Validates tokenizer contracts, vocab sizes, revisions, and teacher-logit metadata.
- `src/kd_contracts.py`: Builds deterministic tokenizer fingerprints.
- `src/training_schedule.py`: Aligns train/validation splits with batch and gradient-accumulation constraints.
- `src/checkpoints.py`: Saves model, tokenizer, scheduler, trainer state, and packing metadata; validates resume compatibility.
- `src/transformers_compat.py`: Resolves attention backend and formats model-loading errors.

## SFT Layer

The `sft/` directory contains the post-KD alignment layer:

- `sft/train_sft.py`: SFT training with optional sequence packing, LoRA/QLoRA paths, and built-in spot evaluations.
- `sft/evaluate.py`: EvalPlus and lm-evaluation-harness orchestration.
- `sft/chat.py`: Local interactive chat wrapper using the tokenizer chat template.

This stage is intentionally separate from KD. KD transfers the teacher's probability structure; SFT teaches the model how to expose that capability in the intended assistant format.