model-a-scratch

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)

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

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.