Modelpublic

namquangstudy/aaie-ddense-gft

sourceHugging Facemitupdated 2mo agoView on Hugging Face

0likes118downloads

Model Card

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 valloss, evaluated every 50 steps — best at end of epoch 2 (valloss 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

python

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.