Model A — From-Scratch Mini Chat Companion
A 3.87M-parameter decoder-only language model built entirely from scratch in pure NumPy — no PyTorch, no JAX, no autograd library. Custom reverse-mode autograd, tokenizer, training loop, KV-cache, and sampler are all hand-written. Optional CuPy backend swap (ZYN_BACKEND=cuda) for GPU training.
Scope is deliberately narrow: short English small-talk / companion replies only. No code generation, no tools, no retrieval, no function-calling — by design.
Architecture (modern-tiny decoder)
| Component | Choice |
|---|---|
| Positions | RoPE (rotary) |
| Norm | RMSNorm, Pre-LN |
| Attention | Grouped-Query Attention (8 query heads, 2 KV heads) + QK-Norm |
| MLP | SwiGLU (2/3 hidden-dim rule) |
| Head | Weight-tied to token embedding |
| Tokenizer | Byte-level BPE, vocab 4096, chat special tokens |
| Layers / d_model / head_dim | 4 / 256 / 32 |
| Context length | 256 |
| Params | 3,869,184 |
| Dtype | float64 (CPU / gradcheck), float32 (GPU) |
Training
Pretrain — DailyDialog (human everyday dialogue), formatted <bos><|user|>…<eos><|assistant|>…<eos>.
- 1.62M tokens · 2000 steps · batch 32 × block 256 · AdamW · cosine LR 3e-4→3e-5
- Val perplexity ≈ 25
Chat fine-tune (SFT) — EmpatheticDialogues (warm, supportive human dialogue) with loss masking (only assistant turns are supervised; user tokens use ignore_index).
- 19.4k dialogues · 800 steps · fresh AdamW · cosine LR 1e-4→1e-5
- Assistant-masked val perplexity ≈ 33 · next-token accuracy 0.325
- Refusal / in-scope check: 3/3 coding prompts answered as chit-chat, no code emitted
Files
| Path | What |
|---|---|
mla/ |
Core library: tensor.py (autograd), model.py, tokenizer.py, kvcache.py, generate.py, chat.py, optim.py, loss.py |
scripts/ |
pretrain.py, finetune_sft.py, build_sft_corpus.py, tokenize_sft.py, evaluate.py |
checkpoints/pretrain_final.npz |
Base pretrained model |
checkpoints/sft_final.npz |
Fine-tuned companion model (use this) |
data/tokenizer/tokenizer.json |
Byte-BPE tokenizer |
tests/ |
Gradcheck, KV-cache equivalence, sampling, chat-runtime tests |
Usage
from mla.checkpoint import load_checkpoint
from mla.tokenizer import Tokenizer
from mla.chat import ChatSession
tok = Tokenizer.load("data/tokenizer/tokenizer.json")
model, _, _ = load_checkpoint("checkpoints/sft_final.npz")
chat = ChatSession(model, tok, temperature=0.8, top_k=40, top_p=0.9)
print(chat.reply("I had a rough day today."))
Inference features: greedy / temperature / top-k / top-p sampling, KV-cache (numerically identical to full forward, verified in tests), multi-turn chat runtime, optional system= persona conditioning.
Limitations
- It is a 3.87M toy. Expect valence errors (may mismatch the emotion of a message), incoherence, and weak multi-turn memory. This is a from-scratch learning artifact, not a production assistant.
- Chat-only: cannot and will not write code, use tools, or do reasoning — out of scope by design.
- Trained on non-commercial data (DailyDialog CC-BY-NC-SA, EmpatheticDialogues CC-BY-NC) → non-commercial use only.
Why it exists
Built atom-by-atom (Karpathy style) to understand every layer of a modern LLM from first principles — each op gradchecked, each stage gated (overfit, KV-cache logit-equivalence, checkpoint resume) before moving on.