Small Llama-style language model
Pretrains a small Llama-style language model, written from scratch in PyTorch, on the TinyStories dataset, then uses it to fit a scaling law and to try the Muon optimizer and a mixture of experts.
- Built the model from the usual Llama parts (RMSNorm, rotary embeddings, grouped-query attention, SwiGLU, a KV cache for sampling), with a flag that turns each feed-forward layer into a Mixtral-style mixture of 8 half-width experts, 2 per token.
- Pretrained a model with 9.7M parameters, not counting embeddings, on 164M tokens in under 6 minutes on an RTX 5070 Ti. It ends at a validation loss of 1.40; its stories have mostly correct grammar but do not always make sense and sometimes repeat themselves.
- Trained five dense models from 0.1M to 9.7M non-embedding parameters on the same 65.5M tokens, with one seed each and no learning-rate tuning per size, and fitted a power law with exponent -0.104 to their validation losses. The largest model sits above the fit, probably short of data at about 7 tokens per parameter.
- Compared Muon and the mixture of experts with an AdamW baseline that has as many active parameters: both end with a lower validation loss but train slower, the mixture of experts 7 to 8 times slower, probably because of my plain Python loop over the experts.
Stack: Python, PyTorch, Hugging Face Tokenizers, NumPy, Matplotlib