| --- |
| license: cc-by-nc-4.0 |
| language: |
| - en |
| library_name: numpy |
| pipeline_tag: text-generation |
| tags: |
| - from-scratch |
| - numpy |
| - small-language-model |
| - chat |
| - companion |
| - rope |
| - rmsnorm |
| - swiglu |
| - gqa |
| --- |
| |
| # 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](https://huggingface.co/datasets/li2017dailydialog/daily_dialog) (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](https://huggingface.co/datasets/Estwld/empathetic_dialogues_llm) (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 |
|
|
| ```python |
| 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. |
|
|
