Ids → embedding → 6 blocks → norm → head → logits

Ids in, logits out. Every step hand-written.

A decoder-only Transformer language model — tokenizer, architecture, training loop, sampler and KV cache — implemented from scratch on PyTorch tensor primitives.

No Hugging Face model classes. No nn.TransformerDecoderLayer, no nn.MultiheadAttention. Attention is written out with matmul and softmax so the mechanism is inspectable, and PyTorch's fused kernel is kept beside it as a reference to check against rather than as the implementation.

The point of the project is not to be a fast Transformer or a good one. It is to be a legible one: every component lives in its own small module, and the test suite checks each against an independently written reference implementation rather than against itself. A model was then genuinely trained on it, and this page reports what that produced — including where it falls short.

Parameters: 6,917,3764,820,224 non-embedding
Vocabulary: 8,192byte-level BPE, trained here
Context: 256tokens, RoPE positions
Depth: 6 × 8blocks × attention heads
Tests: 30719 modules, all passing

The forward pass

One diagram, read top to bottom. Shapes are for the trained configuration at batch B. Position enters only inside attention — there is no learned positional embedding table anywhere in this model.

Token idsoutput of the byte-level BPE tokenizer

(B, 256) int64

Token embeddinga lookup table, nothing more

8192 × 256 → (B, 256, 256)

× 6 decoder blocks

RMSNorm 256

Causal self-attention 8 heads × 32, RoPE on q and k

RMSNorm 256

SwiGLU MLP 256 → 704 → 256

Final RMSNormone last rescale before the projection

(B, 256, 256)

LM headweight-tied to the embedding — one tensor, used twice

256 × 8192

Logitsone score per vocabulary entry, per position

(B, 256, 8192)

The rust spine is the residual stream. Normalisation sits inside each branch, never on the trunk, so the stream is never rescaled and gradients reach the embedding along an unbroken identity path. That is the whole difference between pre-norm and the original 2017 post-norm layout, and it is why this trains at depth without heroics.

What each node is doing

Seven components, each in its own module, each checked against an independent reference implementation in float64.

RMSNorm model/norm.py

y_i = x_i / sqrt( mean_j(x_j²) + ε ) · g_i

LayerNorm with the mean-subtraction and the bias removed — rescale by the root mean square and nothing else. Cheaper, and empirically as good, which is why Llama, Mistral and Qwen all use it.

The sum of squares is accumulated in at least float32 even when activations are bf16: in fp16 a single activation above roughly 256 overflows the moment it is squared. A float64 input is left in float64, so the finite-difference gradient checks stay exact — a hard .float() there would silently narrow them and make the check meaningless.

RoPE — rotary position embeddings model/rope.py

⟨R_mq, R_nk⟩ = ⟨q, R_n−mk⟩

Pairs of channels in the query and key are rotated by an angle proportional to the token's absolute position. The identity above is the entire reason the construction works: rotating by absolute position makes the attention score depend only on the relative offset, so nothing has to be added to the residual stream.

Values are never rotated — only the score carries position. The angle tables are built on CPU in float64 and cast down, because float32 m·w_i loses enough precision at large positions to show up as measurable drift in that identity.

Causal self-attention model/attention.py

scores = q kᵀ / √d_h → mask → softmax → · v

The mask is built from absolute positions: query i may attend to key j only when j ≤ past + i, where past is the number of already-cached tokens. A plain lower-triangular matrix is correct for a full forward pass and silently wrong the moment a KV cache is involved — which is exactly the bug that produces a model that trains beautifully and generates nonsense.

Causality is therefore tested behaviourally rather than by inspecting the mask: perturbing token t must leave every output before it bit-identical, checked at the attention, block and whole-model level.

Grouped-query attention model/attention.py

Setting n_kv_heads below n_heads gives each key/value head to several consecutive query heads, shrinking the KV cache — the dominant memory cost of long-context decoding — by the same factor. The trained model uses plain MHA; the configuration supports the whole range down to multi-query.

Same model width, 1,024-token context, batch 1.
n_kv_heads		Parameters	KV cache
8	MHA	6,917,376	12.00 MiB
4	GQA	6,524,160	6.00 MiB
2	GQA	6,327,552	3.00 MiB
1	MQA	6,229,248	1.50 MiB

The test suite does not merely check that GQA runs. It proves a GQA layer is equal to an MHA layer whose key/value projections are the GQA ones duplicated per group — two parameterisations of the same function.

SwiGLU model/mlp.py

SwiGLU(x) = W_down( SiLU(W_gate x) ⊙ W_up x )

A gated feed-forward block: one branch produces a value, the other a multiplicative gate, so the layer can suppress or pass features per channel rather than applying one fixed nonlinearity. Gating costs a third matrix, so the hidden width defaults to 8/3 · d_model rather than 4 · d_model — 704 here — keeping the parameter count comparable to an ungated MLP of the same width.

Weight tying model/embeddings.py

The LM head and the embedding are one tensor — 2.1M of this model's 6.9M parameters. Tying generally helps at small scale by forcing a token's input and output representations to agree.

It has a consequence worth knowing, and the tests document it: before any training, a tied model's most likely “next” token is the current token, because the residual stream still carries the input embedding. Measuring untrained loss against unshifted targets therefore flatters the model, which is why the initialisation test uses independent targets.

KV cache model/kv_cache.py

Preallocated key/value tensors per layer plus one position cursor shared by all of them. The cursor advances once per forward pass, not once per layer: every layer writes the same slice, and the model commits it after the last block.

Caching changes what is recomputed, never what is computed — so the cached and uncached paths must agree exactly, and that is enforced as a hard gate rather than assumed. Across 432 combinations of attention variant, batch size, prompt length, generated length and sampling strategy, the two paths produce token-for-token identical output. A cache filled entirely with poison values still yields correct logits, because slots beyond the cursor are never read.

What one training step does

Next-token prediction: a window of context_length + 1 tokens is split into inputs and labels shifted by one position. That shift is the objective; the causal mask is what makes it a valid one.

Compute the learning rate from the step index. The schedule is a pure function — linear warmup, then cosine decay — so resuming a run needs only the step number, never a serialised scheduler whose state could drift out of sync with the optimizer.
Run the accumulation micro-batches, each scaled by 1/grad_accum_steps so the accumulated gradient equals the gradient of the mean loss over the full effective batch. Under DDP, all-reduce is suppressed on every micro-step but the last.
Clip by global gradient norm. Globally, across all parameters — clipping per tensor would change the direction of the update, not just its length.
Step AdamW, with weight decay applied only to parameters of rank ≥ 2. Biases and RMSNorm gains are excluded: decaying a normalisation gain fights the layer it belongs to, for a negligible parameter count.

Checkpoints carry the weights, the AdamW moments, the step, the config, and every RNG stream — including the data sampler's own generator. Re-seeding that on resume would replay windows the interrupted run had already trained on, so carrying its state forward is what makes resume exact.

The model that was actually trained

6.9M parameters, 3,000 steps at 8,192 tokens per step — 24.6M tokens, about 8.3 passes over the corpus — in 24.5 minutes of bf16 on one Apple M4 Pro GPU.

Validation loss, measured every 250 steps on held-out documents.
Step	Loss	Perplexity	Bits / byte
250	5.7741	321.86	2.3855
1,000	4.9535	141.67	2.0465
2,000	4.7398	114.41	1.9582
3,000	4.6675	106.43	1.9284

On the full validation split — 97 batches, 397,312 tokens — the final checkpoint scores loss 4.7344, perplexity 113.80, 1.9560 bits per byte. On the training split it scores 3.8616, a 0.87-nat gap after 8.3 epochs: the model has memorised a real amount of its training text. Validation loss was nonetheless still falling monotonically at step 3,000, so this is memorisation alongside generalisation rather than the point where further training starts to hurt.

Perplexity here is not comparable to anything outside this project. It depends on the tokenizer: a coarser vocabulary packs more characters into each token and reports a higher number for identical text. The comparable figure is bits per byte — 1.96, against roughly 1.2 for a good character-level model on English prose. The run is data-bound, not compute-bound: a compute-optimal budget for 6.9M parameters is on the order of 140M tokens, and this corpus has 2.95M.

The corpus is twenty pre-1929 English literary works from Project Gutenberg — Austen, Dickens, the Brontës, Twain, Melville, Doyle, Wells, Stoker, Wilde and others. All are public domain in the United States; Project Gutenberg's own header and footer, which carry its trademark licence, are stripped so only public-domain text reaches disk. The split is by document, not by offset — two whole books are held out — and the tokenizer is trained on the training documents only, so the held-out text influences neither the vocabulary nor the weights.

What it writes

Greedy, temperature, top-k, top-p and a repetition penalty, each a pure function of the logits. Both a cached and an uncached decode path ship, and the second is the reference the first is checked against.

Prompted with “It was a bright cold morning, and ” at temperature 0.8 with top-p 0.95, the trained model continues:

It was a bright cold morning, and ------a big, under the gold head and a bundle of paper. A brown, white eyes, with a black set on the table, and a pink-bed, was a delicate blue, a woman of very white, a table-length, a thin hair, which had hung in her hair, and so her warm hands would have been to me…

Well-formed English words, correct local syntax, consistent punctuation, period-appropriate vocabulary — and no coherence beyond a clause or two. That is what 6.9M parameters trained on 3M tokens buys, and it is shown rather than curated.

Greedy decoding does something worse, and it is worth showing too:

It was a bright cold morning, and 10, 10, 10 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, …

Argmax decoding collapses into a repetition loop. This is the well-known failure mode of greedy sampling, much more pronounced at this scale because the model's distribution is flat enough that the single most likely token is often part of a degenerate cycle. It is exactly why temperature and truncation sampling exist.

What the KV cache is worth

Correctness first, performance second — and the performance result is more interesting than expected.

Greedy decoding, batch 1, 232 new tokens. Median of 9 repeats; both paths verified token-identical before either was timed.
Device	Cached	Uncached	Speedup
CPU, fp32	394 tok/s	54.5 tok/s	7.24×
Apple GPU, bf16	116 tok/s	94 tok/s	1.24×

On CPU the cache does what the textbook says: the speedup rises monotonically from 2.74× at 32 new tokens to 7.24× at 232, which is the difference between O(T) and O(T²). On the Apple GPU it does essentially nothing — both paths sit between 94 and 143 tokens per second regardless of what is being computed. Single-token decoding of a 6.9M-parameter model gives the GPU almost no work per kernel, so wall-clock time is dominated by per-kernel dispatch, and caching removes arithmetic that was never the bottleneck. Re-measuring a 15.7M-parameter model did not bring the effect back, so no claim is made that it returns at scale. This is a fact about one Metal backend at these sizes — not about KV caching, and not about CUDA, which was never available here.

What this is not

The boundary between what was implemented and what was actually executed, stated plainly, because a project like this is only worth anything if its claims are trustworthy.

Not a competitive language model. It is a 6.9M-parameter model trained on 3M tokens of 19th-century fiction. It writes locally plausible English and nothing more. No benchmark evaluation was run and no comparison to any pretrained model is made or implied.
Multi-GPU training is architected, never executed. The DistributedDataParallel integration points exist and are reviewed, but the development machine has one accelerator. No multi-GPU throughput, scaling or convergence number is reported anywhere. Only the degenerate single-process path is covered by tests.
CUDA was never available. Every accelerator figure on this page is Metal. Nothing here should be read as a statement about NVIDIA hardware.
The larger configuration was never trained. A 15.7M-parameter config ships and is validated, and it was benchmarked for throughput with random weights — but no loss, perplexity or sample quality is published for it, because none was measured.
Bitwise resume equivalence is claimed for CPU and float32 only. The reference run was genuinely interrupted and resumed; on the accelerator its validation loss came back 0.0025 nats from the uninterrupted trajectory. That is nondeterministic reduction order in the Metal kernels under bf16, reported as an observation rather than explained away.
Throughput numbers are properties of one laptop. Single-host measurements, taken with warmup excluded and variance reported, on hardware that is named. They do not generalise, and three earlier figures on this project were withdrawn during release preparation when they failed to reproduce under proper methodology.