smolhub-followalong — Llama and Mixtral architectures, from scratch, on a Mac

Take the GPT-2 you understand and turn it into a modern model: RMSNorm, RoPE, grouped-query attention, SwiGLU — then add a Mixture of Experts.

Follows the architectures in SmolHub by Yuvraj Singh (GitHub) — StoryLlama and SmolMixtral — rebuilt for reading and measurement, targeting Apple Silicon.

→ START HERE · → INDEX

---

The whole idea

GPT-2 (2019) and Llama (2023) are the same model with four parts swapped:

	GPT-2	Llama / SmolHub	Why
normalisation	LayerNorm	RMSNorm	Same quality, ~half the work
position	learned embeddings	RoPE	Relative, extrapolates, zero parameters
attention	multi-head	grouped-query	4x smaller KV cache
feed-forward	GELU MLP	SwiGLU	Gated; better at equal parameters

Four small pieces of code, one chapter each. Then the MoE model swaps one more thing: one feed-forward block becomes eight plus a router.

---

python3 -m venv .venv && source .venv/bin/activate
pip install torch numpy tiktoken datasets

python smolhub/model_llama.py     # self-tests: RoPE relativity, KV cache exactness
python smolhub/data.py --tokens 20e6
python smolhub/trainer.py --arch llama --data_dir data/tinystories --dtype bfloat16
python smolhub/inference.py --ckpt ckpt/llama/ckpt.pt \
    --prompt "Once upon a time" --stop_at_eot

A ~15M-parameter model that writes coherent children's stories, in about an hour on an M-series Mac.

---

Prerequisite

Do the companion gpt2-mac repo first, or already know how a GPT-2-style transformer works. Everything here is described as a delta from that.

Chapters

Setup & SmolHub · TinyStories · RMSNorm · RoPE · GQA · SwiGLU · Training StoryLlama · Mixture of Experts · Benchmarking · Evaluation · Is this practical? · Next steps · Troubleshooting

The models self-test

Not decorative — each line is a real correctness assertion:

$ python smolhub/model_llama.py
KV cache matches no-cache (greedy): True
RoPE relative-position property: -2.63799 -2.63799 -2.63799 (should be identical)
  kv_heads=4 (MHA): KV cache at 2048 ctx = 16.8 MB
  kv_heads=2 (GQA): KV cache at 2048 ctx = 8.4 MB
  kv_heads=1 (MQA): KV cache at 2048 ctx = 4.2 MB

Honest findings

The benchmark suite exists to test claims rather than repeat them. Two results worth stating up front:

• The four Llama changes are all worth it. Free or nearly free, and universal for good reason.
• MoE is not worth it on a laptop. ~2.5x the memory to run slower than the equivalent dense model. Its advantage is distributed serving, which a single device cannot realise. Build it to understand it. See Chapter 11.

Requirements

Apple Silicon Mac (M1+), Python 3.10+. torch numpy tiktoken, plus datasets for the TinyStories download. Works on CUDA and CPU with --device.

Licence

MIT. Architectures follow SmolHub (Yuvraj Singh); this is an independent reimplementation for teaching.