Mini LLM

The problem

Most language-model tutorials stop at a notebook. I wanted to understand each part by writing it myself, and then serve the model the way a real one is served: tested, behind an API, with an interface.

The project began as that notebook and was refactored around three questions: which requests the API should refuse before it spends compute, how the notebook’s checkpoints keep loading after the refactor, and what the interface says when the backend or the weights are missing.

What I built

A LLaMA-style decoder-only Transformer with RMSNorm, rotary position embeddings, SwiGLU and causal self-attention, using the GPT-2 tokenizer. Generation uses a KV cache with temperature and top-p sampling. Two-thirds of the parameters are the vocabulary embedding, shared with the output layer.

I trained it in a notebook on WikiText-103: 20,000 steps of 8,192 tokens, 164M tokens in all, with AdamW, a cosine learning-rate schedule and fp16 autocast, on cloud GPUs from Kaggle and Colab. Validation loss fell from 10.92 to 3.64.

The FastAPI service has GET /health, GET /model/info and POST /generate, plus the Swagger page FastAPI generates at /docs. The React console shows the model’s status, takes a prompt with token, temperature and top-p controls, and prints the continuation. Docker Compose runs two containers: the console behind nginx, and FastAPI with the engine imported as a package.

Specification

Parameters

39.1M

Layers

6

Attention heads

8

Context

256 tokens

Hidden size

512

Vocabulary

50,257, GPT-2 tokenizer

Positions

Rotary embeddings (RoPE), precomputed for 512 positions

Feed-forward

SwiGLU, 512 → 768 → 512

Normalisation

RMSNorm

Attention

Causal, 8 heads of 64, PyTorch scaled_dot_product_attention

Dropout

0.1, on the attention output

Decoding

Temperature, then top-p, with a KV cache

Weights

Tied input and output embeddings (64.8M without tying)

Training data

WikiText-103, 20,000 steps, 164M tokens

Optimiser

AdamW, 5e-4 peak, 1,000-step warmup, cosine decay

Key decisions

1. The model knows nothing about the web

   Problem

   Notebook code mixes the model, the experiments and the interface.

   Decision

   Three layers: a React interface, a FastAPI service and an importable PyTorch engine. The frontend never imports model code, and the engine is tested without HTTP.

2. Faster generation with a KV cache

   Problem

   Recomputing attention over the whole prompt for every new token is slow.

   Decision

   Every attention layer returns its keys and values, and each later step feeds only the newest token plus that cache.

   Evidence

   On an Apple M4 CPU the cache is 1.6× faster at 64 new tokens and 2.8× faster at 252 (1.54 s against 4.26 s).

3. Old checkpoints still load

   Problem

   The notebook saved bare state_dict files with no config, and the checkpoint format changed while I was experimenting.

   Decision

   The package keeps the notebook’s parameter names and loads both the raw files and the newer structured checkpoints, which also carry the config, tokenizer name and step.

   Evidence

   The notebook’s own model saved a state_dict, the package loaded it with strict key matching, and both returned identical logits. The final checkpoint has 58 tensors and no missing or unexpected keys.

4. Refuse before spending compute

   Problem

   A prompt plus the requested tokens can exceed the 256-token context, and the rotary table stops at 512 positions.

   Decision

   The API tokenizes the prompt and answers 422 with the reason before the model runs. Sampling settings outside their bounds are rejected the same way, and generate() clamps its own steps for callers that skip the API.

   Evidence

   A 4-token prompt asking for 256 new tokens needs 260 positions, and the refusal took 86 ms.

5. The service starts without weights

   Problem

   Loading the model takes seconds and can fail, but the API and the console should still come up.

   Decision

   /health and /model/info never require a checkpoint. The model loads on first use, a missing file becomes 503 with the fix named, and an unreadable one becomes 500.

   Evidence

   The first /model/info call loads the weights in 1.7–2.1 s; later calls take under 1 ms.

What I measured

Training. Validation loss is cross-entropy over the first 10 validation batches, 81,920 tokens; a model guessing uniformly over 50,257 tokens would score 10.83. The untrained model scores 10.92 (perplexity 55,410), the model after 500 steps 6.00 (404), and the final checkpoint 3.64 (38.1). The 500-step figure comes from re-running the notebook’s recipe on an Apple M4 and stopping at step 800.

Text quality. The output reads like Wikipedia with the meaning taken out, which is what 39M parameters trained on WikiText-103 produce. With the same prompt and seed, temperature 0.3 loops on one sentence, 0.8 writes fluent but invented ones, and 1.5 produces word salad.

Speed. On the M4’s CPU with 4 threads, and the median of three runs, the KV cache makes 64 new tokens 1.6× faster (0.36 s against 0.58 s) and 252 tokens 2.8× faster (1.54 s against 4.26 s), at 164–176 tokens per second. Through the console, an 80-token request took 886 ms from click to text.

Tests. 14 tests in four files pass in 0.58 s on CPU with no checkpoint and no network. They cover the forward pass and the cache shapes, generation and its context clamp, loading both checkpoint formats, and every API route with its bounds and error codes. The frontend has no tests; its production build is the check, and there is no CI yet.

Limits and next steps

A learning model, not a production assistant. It repeats the prompt, ends with a literal end-of-text token and often stops after a sentence or two. The weights aren’t in the repository, so generation needs a checkpoint from the notebook. There is no batching, streaming, authentication or rate limiting, and the engine can run past the trained 256-token context when called directly, though the API refuses it.

Move training from the notebook into a script, save metrics and sample outputs with each checkpoint, publish a small demo checkpoint, add CI for the tests and the frontend build, and add an optional CUDA profile to Docker Compose.