Quintus/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.

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.