# 🌱 Sprout — a tiny chat LLM grown from scratch

Sprout is a **17.3M-parameter** decoder-only transformer trained from scratch on a Mac
(Apple M4 Pro, 3 h 41 min), fine-tuned to chat, and small enough to **run in a web browser**
at ~500 tokens/s with a WebAssembly SIMD engine.

It is the model behind the interactive course **«Своя LLM с нуля» / "Your own LLM from scratch"**
on legost.in, which rebuilds exactly this code chapter by chapter, from counting letters to LoRA:

- 🇬🇧 https://legost.in/en/utilities/micro-llm
- 🇷🇺 https://legost.in/utilities/micro-llm

Talk to Sprout on the course page: everything runs on your device, no server.
Source code: https://github.com/legostin/sprout-llm

```
you> Where do you live?
sprout> I live right inside your web browser.

you> I feel a bit sad today.
sprout> I'm sorry you're feeling this way. I'm here to listen if you want.

you> Tell me a story about a brave little fox.
sprout> Once upon a time, there was a brave little fox. He lived in a big forest with his family.
        The fox had many friends like birds, rabbits, and squirrels. They all played together every day…
```

It is tiny, so it is honest about it: it only speaks English, it can't do arithmetic, and it mixes up
facts. That is the point: every piece of it is small enough to understand.

## The model

| | |
|---|---|
| Parameters | 17,309,056 (3.15M of them in the tied embedding) |
| Architecture | 8 pre-norm blocks, d_model 384, 6 heads × 64, SwiGLU d_ff 1024, context 512 |
| Details | RoPE, RMSNorm, QK-norm, weight tying, no biases, zero-ish init of residual projections |
| Tokenizer | byte-level BPE, 8,192 tokens (4 specials), ≈4.1 characters per token |
| Pre-training | 330M tokens (10,070 steps × 32,768), Muon + AdamW, bf16, `torch.compile`, sequence-length warm-up 128 → 256 → 512, warmup → hold → linear cooldown |
| Result | validation loss 1.554 nats/token ≈ 0.380 nats/letter |
| Chat | supervised fine-tuning on 19,734 conversations, loss only on the assistant's tokens |
| Styles | two LoRA adapters (rank 8, 417,792 trainable numbers): a pirate and a poet |
| Browser | int8 weights with one scale per row (17.5 MB), WebAssembly SIMD matvec ≈14 GMAC/s |

The growth ladder, every model of the course on the same held-out text (nats per letter, lower is better):

| model | chapter | nats/letter |
|---|---|---|
| uniform over 97 symbols | 2 | 4.575 |
| letter frequencies | 2 | 3.084 |
| letter bigram | 1 | 2.364 |
| MLP, 3 letters of context | 4 | 1.479 |
| MLP, 8 letters | 5 | 1.103 |
| MLP on 8 BPE tokens | 6 | 0.809 |
| 1-block transformer | 7 | 0.702 |
| 4-block transformer | 8 | 0.553 |
| **Sprout** | 10 | **0.380** |

Two things that made a real difference on a Mac: **Muon** (at 10M tokens the same model reached
validation loss 2.39 vs 2.71 with a tuned AdamW; orthogonalising same-shaped matrices in one batch keeps
it within 7 % of AdamW's speed) and **`torch.compile`** (≈2× over eager MPS). MLX was not faster
for this size.

## Files

| file | what it does | course chapter |
|---|---|---|
| `download.py` | fetch the raw corpus from Hugging Face (≈6.5 GB) | 9 |
| `tokenizer.py` | byte-level BPE: train, encode, decode | 6 |
| `prepare.py` | clean → dedup → tokenizer → `train.bin` / `val.bin` (1.43B tokens) | 9 |
| `model.py` | the transformer | 7–8 |
| `muon.py` | the Muon optimiser for hidden matrices | 10 |
| `train.py` | pre-training | 10 |
| `build_chats.py` | collect the chat data | 12 |
| `sft.py` | chat fine-tuning | 12 |
| `lora.py` | a style adapter with LoRA, merged back into the weights | 13 |
| `export.py` | int8 weights for the browser | 11 |
| `chat.py` | talk to a checkpoint in the terminal | 12 |
| `snapshots.py` | the small models from the first chapters | 1–6 |
| `course_data.py` | JSON files for the course's widgets | — |
| `web/` | the browser engine: BPE tokenizer (identical to Python), GPT forward with a KV cache, WebAssembly SIMD kernel (`matvec.wat`), sampling, a Worker | 11 |

## Reproduce it

```bash
python -m venv .venv && . .venv/bin/activate
pip install -r requirements.txt

# 1. data
python download.py --out raw                 # ≈6.5 GB
python prepare.py --raw raw --out data       # ≈5 min on 12 cores → 1.43B tokens

# 2. pre-training (3 h 41 min on an M4 Pro, much faster on a CUDA GPU)
python train.py --out runs/base --tokens 330e6 --seq-warmup \
    --sample-at 10,25,50,100,200,400,800 --save-at 200,1000

# 3. chat: grab the ready conversations (see "Materials") or build your own with build_chats.py
python sft.py --base runs/base/ckpt_final.pt --chats data/chats.jsonl --out runs/chat
python chat.py runs/chat/ckpt_final.pt

# 4. a style of your own: 500 conversations in your style is enough
python lora.py --base runs/chat/ckpt_final.pt --chats my-style.jsonl --out runs/my-style

# 5. to the browser
python export.py runs/chat/ckpt_final.pt sprout-chat.bin
```

`train.py` picks CUDA, then MPS, then CPU. On a Mac it runs at ≈25k tokens/s for this size.

## Materials

Everything the model was trained with is public: https://cdn.legost.in/micro-llm/
([MANIFEST.json](https://cdn.legost.in/micro-llm/MANIFEST.json) lists every file with its size and SHA-256).

- `corpus/train.bin`, `corpus/val.bin`: the prepared corpus, uint16 token ids (2.85 GB)
- `tokenizer.json`
- `chats.jsonl`: the 19,734 fine-tuning conversations; `style/pirate.jsonl`, `style/poet.jsonl`
- `checkpoints/sprout-base.pt`, `sprout-chat.pt`, `sprout-pirate.pt`, `sprout-poet.pt` (PyTorch)
- `browser/*.bin`: int8 weights for `web/`, including the small models of the course

About a third of the chat data was written for Sprout's level by language models (Claude) in
18 topic clusters, then mixed with open datasets.

## Data and licences

- [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) (Eldan & Li, 2023), CDLA-Sharing-1.0
- [SimpleStories](https://huggingface.co/datasets/SimpleStories/SimpleStories) (Finke et al., 2025), MIT
- [SODA](https://huggingface.co/datasets/allenai/soda) (Kim et al., 2023), CC BY 4.0
- [everyday-conversations](https://huggingface.co/datasets/HuggingFaceTB/everyday-conversations-llama3.1-2k), Apache-2.0

The code in this repository is MIT-licensed. Muon follows Keller Jordan's
[write-up](https://kellerjordan.github.io/posts/muon/).