89 lines
3.9 KiB
Markdown
89 lines
3.9 KiB
Markdown
|
# 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.
|