AAIE-Distilled Dense Instruct
A ~354M-parameter dense transformer, distilled from Qwen/Qwen2.5-1.5B on general web text
and then instruction-tuned on a general (non-task-specific) instruction dataset. Unlike the
raw AAIE-Distilled Dense checkpoint, this one follows instructions/questions reasonably
rather than just continuing the prompt text — but it has not been specialized for any
particular task (e.g. it was not trained on assignment-feedback data specifically).
This checkpoint exists as the general-purpose midpoint of a longer pipeline: in this project's
experiments, further task-specific fine-tuning on top of this exact checkpoint (rather than
directly on AAIE-Distilled Dense) produced the best result found across the whole investigation
(LLM-judge score 7.50/10 vs. 7.20/10 skipping this stage) — see TECHNICAL_REPORT.md and
SFT_DATA_GENERATION.md for that follow-on experiment. This package ships the general-instruct
checkpoint itself, without that task-specific layer, as a reusable starting point for other
downstream fine-tunes.
Architecture
GQA attention (8 query heads / 2 KV heads) + RoPE + SwiGLU FFN + RMSNorm + tied embeddings,
20 layers, d_model=512, 354M parameters. Tokenizer: same as Qwen/Qwen2.5-0.5B (151,936 vocab).
Data
| Stage | Dataset | Size | Role |
|---|---|---|---|
| Pretraining / distillation | HuggingFaceFW/fineweb-edu (sample-10BT) | ~70B tokens | Base language modeling, distilled against Qwen/Qwen2.5-1.5B. |
| General instruction-tuning | HuggingFaceTB/smoltalk (smol-magpie-ultra config) | 400 examples | Single-turn (first user turn -> first assistant turn) extracted from each conversation. Broad, general-purpose instruction-following signal — not domain-specific. |
Training settings
Distillation (see AAIE-Distilled Dense's model card for full detail): AdamW wd 0.1, LR warmup->3e-4 cosine
decay to 3e-5, 100,000 steps, 131,072 tokens/step, single A100.
General instruction-tuning (this stage, applied on top of the distilled checkpoint):
| Value | |
|---|---|
| Base checkpoint | checkpoints_distill step 100,000 |
| Objective | prompt-masked cross-entropy (loss computed on the completion only) |
| Optimizer | AdamW, weight decay 0.01 |
| LR | constant 3e-6 |
| Batch | 1 example/step, no gradient accumulation |
| Epochs | 2 (~720 steps total over 360 training examples) |
| Grad clip | 1.0 |
| Sequence length | up to 768 tokens (prompt + completion) |
| Train/val split | 90/10, grouped so no data leakage between splits |
| Checkpoint selection | lowest val_loss, evaluated every 50 steps — best at end of epoch 2 (val_loss 2.02, down from an initial 2.28, smooth monotonic improvement with no overfitting observed) |
| Hardware | CPU (small dataset, fast enough without a GPU) |
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "your-username/aaie-ddense-gft" # after pushing, see push_to_hub.py
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, device_map="auto")
prompt = "What are some factors to consider when choosing between two job offers?\n\nResponse:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=150, use_cache=True, do_sample=True, temperature=0.8)
print(tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
trust_remote_code=True is required (custom architecture, see modeling_aaieddenseGFT.py).
use_cache=True (the default) enables real KV-caching for .generate().
Known limitations
- General, not specialized: this checkpoint was tuned on broad instruction data, not on any specific downstream task. If your use case is a specific task (like the assignment-feedback tool this project builds toward), further task-specific fine-tuning on top of this checkpoint is expected to help — that's exactly what this project's experiments found.
- Small SFT set: only 400 general instruction examples. Enough to teach basic instruction-following behavior (confirmed via downstream task performance), not enough for broad instruction-following robustness comparable to production-scale instruct models.
- Fluent but not reliably accurate: like all checkpoints in this project, it can produce confident-sounding but factually wrong output. It was never trained to verify claims, only to produce plausible, on-topic responses.