namquangstudy/aaie-ddense-pretrain

AAIE-Distilled Dense

A ~354M-parameter dense transformer trained on general web text via knowledge distillation against a Qwen/Qwen2.5-1.5B teacher. This is the raw base checkpoint — it has not seen any instruction-tuning or task-specific fine-tuning, so it behaves like a classic base language model: it continues text plausibly, but will not reliably follow instructions phrased as questions or requests (e.g. "Provide feedback on the following...").

If you want a model that follows general instructions, see the companion AAIE-Distilled Dense Instruct model (same base weights, further trained on SmolTalk). If you want the assignment-feedback-specialized model, that's a separate export built on top of both of these stages — see this project's TECHNICAL_REPORT.md.

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

Dataset	Size	Role
HuggingFaceFW/fineweb-edu (sample-10BT subset)	~70B tokens (100,000 steps x 131,072 tokens/step)	General web-text language modeling.
Qwen/Qwen2.5-1.5B (frozen teacher)	—	Distillation signal: loss = 0.5 * CE(labels) + 0.5 * KD(teacher logits, T=2.0).

Training settings

	Value
Base	random initialization
Optimizer	AdamW, weight decay 0.1
LR schedule	warmup 1,000 steps -> 3e-4 peak, cosine decay to 3e-5
Batch	micro-batch 4 x grad-accum 32 x seq-len 1024 = 131,072 tokens/step
Steps	100,000
Grad clip	1.0
Sequence length	1024
Hardware	single A100 (40GB), SLURM cluster

Usage

This is a base model — prompt it as a text continuation, not an instruction:

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "your-username/aaie-ddense-pretrain"  # after pushing, see push_to_hub.py
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, device_map="auto")

inputs = tokenizer("The purpose of a database index is", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=60, 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_aaieddense.py). use_cache=True (the default) enables real KV-caching for .generate() — see GQAAttentionCached in modeling_aaieddense.py.

Known limitations

• Not instruction-tuned: asking it a question or giving it an instruction will often just produce a plausible-sounding continuation of the prompt text rather than an answer/response — this is expected base-model behavior, not a bug.
• Small model, moderate pretraining budget: 354M params on ~70B tokens is well below compute-optimal scale for models trained by major labs; expect weaker world knowledge and reasoning than similarly-sized production models trained on much larger corpora.
• Fluency and factual reliability were not evaluated for this checkpoint specifically — see AAIE-Distilled Dense Instruct and the task-specific model for the evaluations that were run (LLM-judge scoring on IT/CS assignment feedback).