{"slug": "writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in", "title": "Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)", "summary": "Giles Thomas completed building and training a GPT-2 small model from scratch using JAX, achieving a test loss of 3.418784, outperforming both his PyTorch model (3.538161) and the original GPT-2 small (3.499677). The training took 37 hours 15 minutes on an RTX 3090, and the model generated coherent text continuations.", "body_md": "This post is the capstone of\n[the most long-running series on my blog](https://www.gilesthomas.com/llm-from-scratch).\nIn December 2024 (!), I started\nreading [Sebastian Raschka](https://sebastianraschka.com/)'s book\n\"[Build a Large Language Model (from Scratch)](https://www.manning.com/books/build-a-large-language-model-from-scratch)\",\nand worked through it carefully. Being who I am, despite trying to apply a strict \"no side quests\"\npolicy, I found myself zooming off and digging into all kinds of things.\n\nIt's time to wrap it up. I had decided that the endpoint would be to build and train an LLM from scratch just using my notes --\nno reference to the book, no reference to the model code I'd written when following the book.\nAfter an [X/Twitter](https://x.com/gpjt/status/1985434030880293004) poll, I decided to\nuse JAX for that, just to make sure that I really was building it from scratch and\nnot regurgitating bits of PyTorch code like a bad coding LLM spitting out half-digested\nlumps of Stack Overflow.\n\nIn my [last post](https://www.gilesthomas.com/2026/06/llm-from-scratch-34a-building-a-jax-training-loop-for-an-llm-training-run),\nI showed how I built a JAX training script that mirrored what I had built for the original\nPyTorch version of the model. To test it as I went along, I used it to train a really dumb \"LLM\", which instead\nof trying to predict the next token for every token in an input sequence, instead\npredicted the input -- that is, if you fed it\n\n```\nThe fat cat sat on the mat\n```\n\nIt would return the same thing. I called that an A-to-A model.\n\nIn this post, I'll show you how I turned it into a GPT-2 model, and then trained it from scratch on my RTX 3090 (using the parameter counts for the original paper's \"small\" size). What turned out really well with this is that I found a route that meant that almost every component I added made the model better! That's not guaranteed -- sometimes different aspects of an AI model depend on each other, so adding A without also adding B makes things worse. But (admittedly with a bit of backtracking in places) I was able to find a route that shows a nice clear progression.\n\nThe final training run took 37 hours 15 minutes -- compared to 40 hours, 38 minutes for\n[an equivalent PyTorch model](https://www.gilesthomas.com/2026/04/llm-from-scratch-32k-interventions-training-our-best-model-locally-gradient-accumulation).\nThat is despite it being full-fat 32-bit -- the PyTorch one was using Automatic\nMixed Precision (AMP), which allowed it to use 16-bit calculations in places where it would\nbe relatively harmless in terms of loss.\n\nWhen asked to continue \"Every effort moves you\", it came back with a decent response:\n\n```\nEvery effort moves you closer to your goals, but if you are unsure of what it takes, you don’t\n```\n\nThe model got 3.418784 loss on my held-back test dataset, as compared to my PyTorch\nmodel's 3.538161, and even more impressively, it was better than the original GPT-2 small's\nresult of 3.499677 on the same dataset! However, just as [I found previously](https://www.gilesthomas.com/2026/04/llm-from-scratch-32l-interventions-instruction-fine-tuning-tests),\nthe OpenAI weights still beat mine consistently in instruction fine-tuning challenges.\n\nLet's get started.\n\nAt the end of the last post, we had a solid training loop, using all of the tricks I'd picked up with my PyTorch code. The A-to-A model we were training with it looked like this:\n\n``` python\nfrom flax import nnx\n\nclass GPTModel(nnx.Module):\n\n    def __init__(\n        self,\n        vocab_size, context_length,\n        d_emb,\n        n_heads, n_layers,\n        qkv_bias,\n        drop_rate,\n        rngs,\n    ):\n        self.token_embedding = nnx.Embed(\n            num_embeddings=vocab_size,\n            features=d_emb,\n            rngs=rngs,\n        )\n\n        self.output_head = nnx.Linear(\n            in_features=d_emb,\n            out_features=vocab_size,\n            use_bias=False,\n            rngs=rngs,\n        )\n\n    def __call__(self, xs):\n        input_embeddings = self.token_embedding(xs)\n\n        return self.output_head(input_embeddings)\n```\n\nThat was based on my preferred model of [how LLMs work](https://www.gilesthomas.com/2025/09/how-do-llms-work), where at the top level\nfor a model, we feed in a sequence of token IDs, then:\n\nThe A-to-A model basically skipped the second step completely: it would project to embedding space, then immediately project back to vocab space -- and after training, it was pretty good at mapping a sequence to itself.\n\nOne interesting question is, if we train the same code, but this time try to get it to make next-token predictions, how good will it be at that? Obviously it can't be as good as a full LLM. But there are correlations between tokens; full stops will generally be followed by spaces, adjectives will normally be followed by other adjectives or nouns (at least in English), and so on. It would be kind of like the predictive text systems on a phone, where (at least until recently) it would just use the last word you entered to generate a list of possible next words to select from.\n\nOld-school natural language processing has a name for this: bigrams. The idea is that you can work out statistically what the most common two-word pairs are, which allows you to make a guess at a next word from a single one. (There are also trigrams, where you look at the last two words when predicting the next, then 4-grams, 5-grams, and so on.) You'd build up a full probability table -- for every word in your vocab, you'd have the probability of every word coming next.\n\nSo maybe even with that minimal model, we could get it to learn something similar to a set of token-level (rather than word-level) bigrams, which would then get the loss down. Obviously it wouldn't be as good as a full bigram table -- for our GPT-2 vocab size of 50,257, that would need parameters -- but perhaps it could approximate one. (For comparison, the model we're using has just an embedding table and an output head, each mapping between 50,257 dimensions and 768, so that's parameters -- about 3% of the full table.)\n\nAn uninitialised model would (hopefully) have a loss of about\n10.82, implying a\n[perplexity](https://www.gilesthomas.com/2025/10/llm-from-scratch-21-perplexed-by-perplexity) equal to\nthe vocab size. If we can train our dumb model to get better loss than that, then\nwe'd have the beginnings of an LLM.\n\nThat was a simple test to run. In my training code, I had a dataset class that looked like this:\n\n``` python\nclass BigTrainDataset:\n\n    def __init__(self, all_tokens, seq_length, microbatch_size):\n        self.xs = all_tokens[:-1].reshape(-1, microbatch_size, seq_length)\n        self.ys = all_tokens[:-1].reshape(-1, microbatch_size, seq_length)\n\n    def __getitem__(self, ix):\n        return self.xs[ix], self.ys[ix]\n\n    def __len__(self):\n        return self.xs.shape[0]\n```\n\nThat is, the inputs, the `xs`\n\n, were the same as the targets, the `ys`\n\n. If we fed it\n\n```\nThe fat cat sat on the mat\n```\n\n...then we'd be training it to output exactly the same thing.\n\nThe modified version for a real LLM would involve feeding it something like this:\n\n```\nThe fat cat sat on the\n```\n\n...and targeting this:\n\n```\n fat cat sat on the mat\n```\n\nThat's a simple change -- that `__init__`\n\nmethod became this:\n\n```\n        self.xs = all_tokens[:-1].reshape(-1, microbatch_size, seq_length)\n        self.ys = all_tokens[1:].reshape(-1, microbatch_size, seq_length)\n```\n\nI did that, and kicked it off to train on the 92,209,152 tokens that I was (somewhat arbitrarily) using in the last post to test my training loop. The loss chart looked like this:\n\nThat was pretty promising! Loss came down from roughly 10.82 down to a fairly stable 6 or so by global step 768, and seemed to flatten out there. It's possible that further training could have got it down a bit more, but I decided (again, somewhat arbitrarily) to use the average train loss in the checkpoint period ending at step 937 as my starting point. If we could make changes that reduced that, then we'd be moving forward. For this model, that value was 5.909.\n\nSo, what were the changes we needed to make to change our bigram-style model to a real, if small, LLM?\n\nAdapting from my [how LLMs work](https://www.gilesthomas.com/2025/09/how-do-llms-work) post, a GPT-2-style LLM\nlooks like this. We receive our sequence of token IDs, and then:\n\nInside the Transformers blocks, we:\n\nSo that gave me the checklist; looking at it, the most tempting next step was layer normalisation (henceforth LayerNorm). It's used at the end of the core loop, and then twice in the Transformers blocks.\n\nWhat would happen if we coded it up, and then added it to the core only?\n\nThe purpose of [LayerNorm](https://www.gilesthomas.com/2025/07/llm-from-scratch-16-layer-normalisation) is to stabilise\ntraining. We constrain the values flowing through our model so that they have certain\nstatistical properties that tend to make the whole thing more trainable.\n\nThat would mean that if it did help with this model -- placed in between the embedding layer at the start, and the output head at the end -- then we'd hope for loss to go down faster, and ideally finish at a lower level.\n\nTime to code it up!\n\n[NNX has its own LayerNorm implementation](https://flax.readthedocs.io/en/stable/api_reference/flax.nnx/nn/normalization.html#flax.nnx.LayerNorm),\nof course (as does [PyTorch](https://docs.pytorch.org/docs/2.12/generated/torch.nn.LayerNorm.html)),\nbut in the book, we implement it ourselves, and that felt like the correct path to\ntake.\n\nFirstly, I implemented a dummy version:\n\n``` python\nclass LayerNorm(nnx.Module):\n\n    def __init__(self):\n        ...\n\n    def __call__(self, xs):\n        return xs\n```\n\n...and updated the core `GPTModel`\n\nto create and call one:\n\n``` python\nclass GPTModel(nnx.Module):\n    def __init__(\n        ...\n    ):\n        self.token_embedding = nnx.Embed(\n                ...\n        )\n\n        self.output_norm = LayerNorm()\n\n        self.output_head = nnx.Linear(\n                ...\n        )\n\n    def __call__(self, xs):\n        input_embeddings = self.token_embedding(xs)\n\n        normalised = self.output_norm(input_embeddings)\n\n        return self.output_head(normalised)\n```\n\nAnd kicked off a training run for a few seconds just to make sure that it hadn't broken anything and that loss dropped -- being my first NNX module-inside-a-module, I worried that there might have been something non-intuitive that I had to do to get it to work. But everything seemed good -- loss was dropping, no errors.\n\nSo, following [the notes I made when I first learned about LayerNorm](https://www.gilesthomas.com/2025/07/llm-from-scratch-16-layer-normalisation),\nI needed to make the values flowing through centred around zero by subtracting their\nmean, and then scale them to have a variance of one by dividing by the standard\ndeviation (details in those notes).\n\nThe shape of the `xs`\n\nI had coming into my `LayerNorm`\n\nclass's `_call`\n\nwas this:\n\n```\nxs.shape=(6, 1024, 768)\n```\n\nThat was `(batch_size, seq_len, d_emb)`\n\n. So we needed to do those operations strictly on the last axis,\nmanipulating each embedding independently.\n\nJAX has a [ std function](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.std.html)\nand a\n\n`mean`\n\none`axis`\n\nparameter. The `Array`\n\nobject repackaged those as\nmethods, which was convenient, so I did a first cut test like this:\n\n``` python\nclass LayerNorm(nnx.Module):\n\n    def __init__(self):\n        ...\n\n    def __call__(self, xs):\n        jax.debug.print(f\"{xs.shape=}\")\n\n        means = xs.mean(axis=-1)\n        jax.debug.print(f\"{means.shape=}\")\n\n        stds = xs.std(axis=-1)\n        jax.debug.print(f\"{stds.shape=}\")\n\n        return xs\n```\n\nThat printed out these results:\n\n```\nxs.shape=(6, 1024, 768)\nmeans.shape=(6, 1024)\nstds.shape=(6, 1024)\n```\n\n...which looked plausible; one number for each embedding vector. Could we broadcast them across the array?\n\n``` python\nclass LayerNorm(nnx.Module):\n\n    def __init__(self):\n        ...\n\n    def __call__(self, xs):\n        jax.debug.print(f\"{xs.shape=}\")\n\n        means = xs.mean(axis=-1)\n        jax.debug.print(f\"{means.shape=}\")\n\n        stds = xs.std(axis=-1)\n        jax.debug.print(f\"{stds.shape=}\")\n\n        normalized = (xs - means) / stds\n        jax.debug.print(f\"{normalized.shape=}\")\n\n        return normalized\n```\n\nThis blew up:\n\n```\nValueError: Incompatible shapes for broadcasting: shapes=[(6, 1024, 768), (6, 1024)]\n```\n\nFair enough. But `mean`\n\nand `std`\n\nhave a `keepdims`\n\nkwarg that looked like it would help:\n\n``` python\nclass LayerNorm(nnx.Module):\n\n    def __init__(self):\n        ...\n\n    def __call__(self, xs):\n        jax.debug.print(f\"{xs.shape=}\")\n\n        means = xs.mean(axis=-1, keepdims=True)\n        jax.debug.print(f\"{means.shape=}\")\n\n        stds = xs.std(axis=-1, keepdims=True)\n        jax.debug.print(f\"{stds.shape=}\")\n\n        normalized = (xs - means) / stds\n        jax.debug.print(f\"{normalized.shape=}\")\n\n        return normalized\n```\n\n...and it did!\n\n```\nxs.shape=(6, 1024, 768)\nmeans.shape=(6, 1024, 1)\nstds.shape=(6, 1024, 1)\nnormalized.shape=(6, 1024, 768)\n```\n\nExcellent. So the next step was to see if that would work even slightly. Interestingly\nloss started off a bit higher at 11.29 after the first global step -- so adding in the\nLayerNorm had actually made the model *worse* than it was -- but it seemed to be falling\nrapidly. Things weren't totally broken, at least.\n\nBut there was more to LayerNorm than just zeroing the mean and scaling to the variance; we also needed to scale them up by a learnable amount, and then shift/bias them by adding on a different trainable amount. More precisely, both of those trainable amounts were different for each of the (in this case) 768 embedding dimensions.\n\nWe needed two learnable vectors of length `d_emb`\n\n. I hadn't noted it down at the\ntime but I figured (as it turned out, correctly) that a sensible starting point for\nthose values would be all-zero for the bias, and all-one for the scale.\n\nFrom [this help page](https://flax.readthedocs.io/en/stable/key_concepts.html#traced-vs-static-data),\nthe way you create a trainable array associated with an NNX module is this:\n\n```\nnnx.Param(jax.random.normal(rngs.param(), (dim, dim)))\n```\n\nThat code created a random vector, rather than the zeros/ones we needed, and we'd need to get\nthe dimensions right. Because of the \"Incompatible shapes for broadcasting\" error I'd\njust had, I was feeling a bit paranoid about the latter, so I chose a shape of `(1, 1, d_emb)`\n\n, and\nwrote this:\n\n``` python\nclass LayerNorm(nnx.Module):\n\n    def __init__(self, d_emb):\n        self.scale = nnx.Param(jnp.ones((1, 1, d_emb)))\n        self.bias = nnx.Param(jnp.zeros((1, 1, d_emb)))\n\n    def __call__(self, xs):\n        jax.debug.print(f\"{xs.shape=}\")\n\n        means = xs.mean(axis=-1, keepdims=True)\n        jax.debug.print(f\"{means.shape=}\")\n\n        stds = xs.std(axis=-1, keepdims=True)\n        jax.debug.print(f\"{stds.shape=}\")\n\n        normalized = (xs - means) / stds\n        jax.debug.print(f\"{normalized.shape=}\")\n\n        scaled_and_biased = (normalized * self.scale) + self.bias\n\n        return scaled_and_biased\n```\n\nThat looked pretty plausible, though in retrospect I think I was being overly cautious and didn't need the leading two axes for the scale and bias.\n\nThe only thing I was unsure about was whether the `nnx.Param`\n\nwrappers\nI had put in were really making those arrays trainable. I put some code in to print\nthem out and kicked off a run for a few minutes, and confirmed that they were\nchanging in ways that seem plausible -- small non-zero bias, scale close\nto but not equal to one. That was all good!\n\nNext, I spotted one issue. What if one of the standard deviations was zero? That would lead to a divide-by-zero error here:\n\n```\nnormalized = (xs - means) / stds\n```\n\nNow, the standard deviation, if it's not zero, has to be positive --\nso adding on a small value would fix that 1:\n\n```\n        normalized = (xs - means) / (stds + 1e-5)\n```\n\nWith that in place, I felt that it was ready to go. Time to do a full training run!\n\nI kicked that off, and it completed with this output:\n\n```\n2026-06-20 19:08:17.189721 Tokens seen: 92,209,152\n2026-06-20 19:08:17.189724 Throughput: 95,383 tokens/second\n2026-06-20 19:08:17.189734 Final train loss: 5.736\n2026-06-20 19:08:17.189737 Done\n```\n\nLoss looked like this:\n\nLet's look at the results for the previous run without LayerNorm for comparison:\n\nYou can see that the new run, the first one, drops faster. It's harder to see from the chart, but it also finished up with a lower training loss at 937 (my relatively arbitrary metric): 5.734 rather than 5.909.\n\nThat was interesting! The new model was basically doing the same thing -- predicting the next token based only on the \"current\" token, but loss was lower. My take is that if we had trained the non-LayerNorm model for longer, it might have managed to eventually grind out a better loss. But LayerNorm was doing its job -- it was stabilising training, and as a result we converged faster.\n\nThat was a win! I decided to run it through my old smoke test from the PyTorch training runs, and see how it completed \"Every effort moves you\":\n\n```\nEvery effort moves you can be a few years.\n-year-year-year-year-year-year-\n```\n\nIt was kind of impressive that it managed to finish the first line before it got stuck in a loop -- but it was understandable that we couldn't expect anything good yet. Each predicted token was based entirely on the token before it.\n\nWhat next? Back to our checklist:\n\nInside the Transformers blocks, we:\n\nSo, at this stage, for each input token we were predicting the next one based on the input token only -- like I said earlier, we were doing a somewhat roundabout way of building an approximation of a table of bigram probabilities. What would happen if we started paying attention to the tokens to the left? And what would be the simplest, dumbest way to do that?\n\nThe real LLM has multiple layers of multi-head attention, each one also having a feed-forward network, some LayerNorms, and some shortcut connections.\n\nSingle-head attention is easier to code, but even on its own, you'd expect it to be able to add some value. Each token would get at least some information from the ones to the left. And one layer, likewise, you'd expect might help a bit. I suspected that it wouldn't work on its own -- I expected I'd need shortcut connections too -- but decided to start with attention on its own.\n\nI modified the main class to have a single \"Transformers\" layer:\n\n``` python\nclass GPTModel(nnx.Module):\n\n    def __init__(\n        ...\n    ):\n        self.token_embedding = nnx.Embed(\n            ...\n        )\n\n        self.transformers_layer = TransformersLayer(d_emb, qkv_bias, rngs)\n\n        self.output_norm = LayerNorm(d_emb)\n\n        self.output_head = nnx.Linear(\n            ...\n        )\n\n    def __call__(self, xs):\n        input_embeddings = self.token_embedding(xs)\n\n        transformed = self.transformers_layer(input_embeddings)\n\n        normalized = self.output_norm(transformed)\n\n        return self.output_head(normalized)\n```\n\n...where that layer was actually just single-head attention:\n\n``` python\nclass TransformersLayer(nnx.Module):\n\n    def __init__(self, d_emb, qkv_bias, rngs):\n        self.attention = Attention(d_emb, qkv_bias, rngs)\n\n    def __call__(self, xs):\n        return self.attention(xs)\n```\n\nNext, it was time for the `Attention`\n\nclass.\nI'm not going to write yet another attention explainer -- I think my [\"How do LLMs work?\"](https://www.gilesthomas.com/2025/09/how-do-llms-work)\none does a decent job of that, and [\"The 'why' of attention, or: attention heads are dumb\"](https://www.gilesthomas.com/2025/05/llm-from-scratch-13-taking-stock-part-1-attention-heads-are-dumb)\nworks well too. So in the next bit I'll assume that you understand the basics.\n\nMy first cut was basically just the maths (up to the causal mask) to get the attention scores:\n\n``` python\nclass Attention(nnx.Module):\n\n    def __init__(self, d_emb, qkv_bias, rngs):\n        self.d_emb = d_emb\n\n        self.W_q = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n        self.W_k = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n        self.W_v = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n\n    def __call__(self, xs):\n        Q = self.W_q(xs)\n        K = self.W_k(xs)\n        V = self.W_v(xs)\n\n        omega = Q @ K.T\n\n        omega /= jnp.sqrt(self.d_emb)\n\n        causal_omega = jnp.tril(omega)\n```\n\nIt did the projections into query, key and value space, worked out the attention scores with the array multiplication, normalised it by dividing by the square root of the number of dimensions in the Q-K embedding space, and then zeroed out the scores where a token was attending to tokens in its \"future\".\n\nThere were a couple of problems, though. Firstly, that wouldn't work if we were working with batches, and secondly, zeroing out the non-causal scores wasn't quite correct.\n\nThe batches first. Our incoming `xs`\n\nhere would have the shape\n`(batch_length, seq_len, d_emb)`\n\n. After the projections to the Q-K embedding space,\nboth `Q`\n\nand `K`\n\nwould also be shaped `(batch_length, seq_len, d_emb)`\n\n. Now, the\n`.T`\n\nproperty on the JAX array class just reverses the axes, so the code above\nwould give us `K.T`\n\nwith the shape `(d_emb, seq_len, batch_length)`\n\n.\n\nThat would break! Matrix multiplication in JAX expects all but the last two\naxes to represent batches, so we actually wanted `K.T`\n\nto have the shape\n``(batch_length, d_emb, seq_len)`\n\n. That meant that what we actually\nwanted was to just transpose the last two axes.\n\nThe JAX [ transpose function](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.transpose.html) takes\nan\n\n`axes`\n\nparameter that allows you to specify the specific re-ordering of the input\naxes that you want. So I could rewrite the code like this:\n\n``` python\n    def __call__(self, xs):\n        Q = self.W_q(xs)\n        K = self.W_k(xs)\n        V = self.W_v(xs)\n\n        omega = Q @ jnp.transpose(K, axes=(0, 2, 1))\n\n        omega /= jnp.sqrt(self.d_emb)\n\n        causal_omega = jnp.tril(omega)\n```\n\nAs `Q`\n\nwould have the shape `(batch_length, seq_len, d_emb)`\n\n, and the transposed\nversion of `K`\n\nwould be `(batch_length, d_emb, seq_len)`\n\n, they'd be compatible for\nmatrix multiplication and give us a result that was `(batch_length, seq_len, seq_len)`\n\n-- just what we wanted for attention scores.\n\nThe next step was to fix the causal mask. The next step in this\nattention mechanism was going to be running the causal attention scores\nin `causal_omega`\n\nthrough softmax over the last dimension, to convert them into\nattention weights. Now, our current code was zeroing out unwanted acausal scores,\nbut a zero still contributes to softmax. If you want a particular value to come out of\nsoftmax guaranteed to be zero, you need to set it to minus infinity.\n\nI decided that the easiest way to do this was to create a causal mask -- a boolean\narray that matched the size of `omega`\n\n, but was full of `True`\n\ns:\n\n```\ncausal_mask = jnp.ones_like(omega, dtype=bool)\n```\n\nThen I could zero out (well, \"false out\") the cells in the mask related to unwanted future-facing scores, just like I was previously doing on the scores:\n\n```\ncausal_mask = jnp.tril(causal_mask)\n```\n\n...and then I could apply that mask to omega with [ jnp.where](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.where.html),\ntelling it to create a new array, taking the value from\n\n`omega`\n\nwhere the mask\nhad `True`\n\n, and `-jnp.inf`\n\nin places where it had `False`\n\n.\n\n```\ncausal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n```\n\nThat seemed solid, so I just needed to run the result through\n[ jax.nn.softmax](https://docs.jax.dev/en/latest/_autosummary/jax.nn.softmax.html),\nspecifying that the last dimension was the one where it should apply the function,\nand that would give me the attention weights:\n\n```\nattention_weights = jax.nn.softmax(causal_omega, axis=-1)\n```\n\nFinally, I just needed to use those attention weights to get the attention output\nby mixing in appropriate portions of the projection of the inputs into value space,\n`V`\n\n:\n\n```\nreturn attention_weights @ V\n```\n\nAs `attention_weights`\n\nwas shaped `(batch_length, seq_len, seq_len)`\n\n, and `V`\n\n(like\n`Q`\n\nand `K`\n\n) was shaped `(batch_length, seq_len, d_emb)`\n\n, the batch axes were at the\nstart where they belonged, and the matrix multiplication would work and return something\nshaped `(batch_length, seq_len, d_emb)`\n\n.\n\nWith that, we were done! The final single-head attention class looked like this:\n\n``` python\nclass Attention(nnx.Module):\n\n    def __init__(self, d_emb, qkv_bias, rngs):\n        self.d_emb = d_emb\n\n        self.W_q = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n        self.W_k = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n        self.W_v = nnx.Linear(d_emb, d_emb, use_bias=qkv_bias, rngs=rngs)\n\n    def __call__(self, xs):\n        Q = self.W_q(xs)\n        K = self.W_k(xs)\n        V = self.W_v(xs)\n\n        omega = Q @ jnp.transpose(K, axes=(0, 2, 1))\n\n        omega /= jnp.sqrt(self.d_emb)\n\n        causal_mask = jnp.ones_like(omega, dtype=bool)\n        causal_mask = jnp.tril(causal_mask)\n\n        causal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n\n        attention_weights = jax.nn.softmax(causal_omega, axis=-1)\n\n        return attention_weights @ V\n```\n\nI kicked off a training run with that, and it did work, in that loss went down over the course of the run -- but at the end of the run, the loss at step 937 was 5.934 -- significantly above the 5.734 I got on the previous run, with no attention.\n\nBut that made sense! As I'd said earlier, I suspected that this wouldn't help if we had no shortcut connection.\n\nIntuitively, if you want to work out what token should be at position , on average the\nmost important other token you need to know about is probably whichever one is at position .\nKnowing about the tokens at , , and so on, could well be helpful -- maybe\n*very* helpful -- but not at the cost of not knowing about the one at .\n\nNow, single attention heads are just simple pattern-matchers. They can't learn complex rules, it's only by working together -- \"horizontally\", in multi-head attention or \"vertically\" across multiple layers -- that they can do complex things.\n\nWhat we were asking this head to do was to learn some way of gathering information about previous tokens, and also to keep the knowledge about the \"current\" one. That's a tall order for a dumb attention head!\n\nIn my mind, this is a large part of the benefit of shortcut connections.\nThey are often presented as a way to make sure that during training,\ngradients flow smoothly from the output end of the model to the earlier layers. But\nI prefer to think of them as preserving the original embeddings, so that each layer doesn't\ncompletely replace what came into it, but instead does something closer to adding on\nits own notes -- like [scholars adding commentary to a core text in the Talmud](https://www.gilesthomas.com/2025/08/llm-from-scratch-18-residuals-shortcut-connections-and-the-talmud).\n\nIn the training run above, the attention head was trying to learn how to preserve the meaning of the embedding it was working on, while also merging in information from earlier ones. If we added a shortcut connection, then it would only have to do the second of those two jobs.\n\nThe code was simple: I updated the `TransformersLayer`\n\nmodule to do a shortcut\nconnection:\n\n``` python\nclass TransformersLayer(nnx.Module):\n\n    def __init__(self, d_emb, qkv_bias, rngs):\n        self.attention = Attention(d_emb, qkv_bias, rngs)\n\n    def __call__(self, xs):\n        shortcut = xs\n        att = self.attention(xs)\n        return shortcut + att\n```\n\nI kicked off a training run, and at the end it printed this:\n\n```\n2026-06-23 03:51:18.086097 Tokens seen: 92,209,152\n2026-06-23 03:51:18.086099 Throughput: 90,442 tokens/second\n2026-06-23 03:51:18.086108 Final train loss: 5.570\n2026-06-23 03:51:18.086121 Done\n```\n\nThe loss chart looked like this:\n\nAnd, importantly, that training loss at step 937 which I was using as a metric was 5.553 -- a decent improvement over the previous best of 5.734. Even a dumb single attention head was able to do something useful, if it had a shortcut connection.\n\nI decided to run another qualitative smoke test:\n\n```\nEvery effort moves you can be able to get to get to get to get to get a lot of the way to get\n```\n\nI mean, it was repetitive, but it was actually getting noticeably closer to making sense!\n\nSo that was excellent news. What next? Our checklist looked like this:\n\nInside the Transformers blocks, we:\n\nNow, our single attention layer was lacking something. Without position embeddings, that layer has no idea what order the tokens before the one it's looking at come in. If it's considering the \" cat\" in\n\n```\nThe fat cat\n```\n\n...it doesn't know if it's looking at \"The fat cat\" or \"fat The cat\".\n\nPosition embeddings are simple, and might help, so that was the next step.\n\nThese were trivial to add. We had this core code:\n\n``` python\nclass GPTModel(nnx.Module):\n\n    def __init__(\n        ...\n    ):\n        self.token_embedding = nnx.Embed(\n            ...\n        )\n\n        self.transformers_layer = TransformersLayer(d_emb, qkv_bias, rngs)\n\n        self.output_norm = LayerNorm(d_emb)\n\n        self.output_head = nnx.Linear(\n            ...\n        )\n\n    def __call__(self, xs):\n        input_embeddings = self.token_embedding(xs)\n\n        transformed = self.transformers_layer(input_embeddings)\n\n        normalized = self.output_norm(transformed)\n\n        return self.output_head(normalized)\n```\n\nSo I just added a position encoding module in `__init__`\n\n:\n\n```\n        self.position_embedding = nnx.Embed(\n            num_embeddings=context_length,\n            features=d_emb,\n            rngs=rngs,\n        )\n```\n\n...and mixed it in with the token embeddings to create new, improved `input_embeddings`\n\nto be used\nin our \"Transformers\" layer:\n\n```\n        token_embeddings = self.token_embedding(xs)\n        b, n = xs.shape\n        position_embeddings = self.position_embedding(jnp.arange(n))\n        input_embeddings = token_embeddings + position_embeddings\n```\n\nI kicked off a training run with that:\n\n```\n2026-06-23 04:44:44.759768 Tokens seen: 92,209,152\n2026-06-23 04:44:44.759771 Throughput: 88,618 tokens/second\n2026-06-23 04:44:44.759779 Final train loss: 5.386\n2026-06-23 04:44:44.759781 Done\n```\n\nPretty hard to distinguish from the previous one, but the metric I was tracking, that loss at step 937, had improved again! We were down to 5.354 from 5.553 :-)\n\nA quick qualitative smoke test didn't show that improvement, though:\n\n```\nEvery effort moves you can be able to get to get to get back to get back to get back to get back to\n```\n\nPretty much indistinguishable to the previous one. But still, Loss Number Went Down, and that's what was important at this stage.\n\nIt was time to try the next step. From the checklist:\n\nInside the Transformers blocks, we:\n\nWe had only one attention head right now. Individually,\n[attention heads are dumb](https://www.gilesthomas.com/2025/05/llm-from-scratch-13-taking-stock-part-1-attention-heads-are-dumb),\nso switching to multi-head attention seemed like a good thread to pull.\n\nAt this point, my single-head attention code looked like this:\n\n```\n        Q = self.W_q(xs)\n        K = self.W_k(xs)\n        V = self.W_v(xs)\n\n        omega = Q @ jnp.transpose(K, axes=(0, 2, 1))\n\n        omega /= jnp.sqrt(self.d_emb)\n\n        causal_mask = jnp.ones_like(omega, dtype=bool)\n        causal_mask = jnp.tril(causal_mask)\n\n        causal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n\n        attention_weights = jax.nn.softmax(causal_omega, axis=-1)\n\n        return attention_weights @ V\n```\n\nI decided to re-implement multi-head attention (which I'll call MHA from here onwards) from first principles rather than working strictly from my notes, and then to come back and check it.\n\nIf you're looking at your browser's scrollbar with horror (\"\n\nstillonly 50%?!\") and really don't want to read a full derivation of MHA, you can[skip straight to the first complete version of the code].\n\nThe point of MHA is that we're running multiple copies of the calculation above in\nparallel -- let's pin down the name of the number of copies as `n_heads`\n\n. Now,\nwe could naively implement it just by spinning off `n_heads`\n\nthreads and running the\nexisting code in each, but that wouldn't really take advantage of the GPU's inherent\nparallelism.\n\nI felt that we could rely on the fact that JAX's matrix multiplications treat all but the last two dimensions as \"batches\". For example, if you have two arrays with shapes:\n\n```\n(a, b, c, ..., l, m, n)\n```\n\nand\n\n```\n(a, b, c, ..., l, n, p)\n```\n\n...then you can multiply them. A matrix multiplied by a one will be , so you'll get something that is\n\n```\n(a, b, c, ..., l, m, p)\n```\n\nThe other dimensions (so long as they match) will essentially act as an batch.\n\nNow, right now we were just using a single batch dimension. Let's look at the core\nmultiplication in the attention mechanism, which works out `omega`\n\n, the attention scores.\nI had this:\n\n```\nomega = Q @ jnp.transpose(K, axes=(0, 2, 1))\n```\n\nBreaking that apart into two steps:\n\n```\nK_transpose = jnp.transpose(K, axes=(0, 2, 1))\nomega = Q @ K_transpose\n```\n\nWe got `K`\n\nfrom this line:\n\n```\nK = self.W_k(xs)\n```\n\nLet's look at the shapes here.\n\n`xs`\n\nis our input embeddings for this layer; its\nshape is `(batch_size, seq_len, d_emb)`\n\n. Projecting it through `W_k`\n\n, which is\nshaped `(d_emb, d_emb)`\n\ngives us a shape for `K`\n\nof `(batch_size, seq_len, d_emb)`\n\nagain.\n\n`Q`\n\n, being a projection of `xs`\n\nthrough `W_q`\n\n, which is the same shape as `W_k`\n\n, will have\nthe same shape as `K`\n\n.\n\nNow, that means that `K_transpose`\n\nis `(batch_size, d_emb, seq_len)`\n\n, and the calculation\n\n```\nomega = Q @ K_transpose\n```\n\n...is doing a batched matrix multiplication getting us the `omega`\n\nthat we want,\nshaped `(batch_size, seq_len, seq_len)`\n\n.\n\nBut as I said above, there's no need to stop with just one batch dimension. Let's say\nthat we have `n_heads`\n\nheads, and that they each work with embeddings sized `d_head`\n\n.\nImagine that we've already somehow done multiple projections into the key and query\nspaces for each of our `n_heads`\n\nheads, and that the results have somehow been put\ninto arrays such that `Q`\n\nand `K`\n\nare shaped `(batch_size, n_heads, seq_len, d_head)`\n\n-- that is, we've gained an extra axis that keeps the projections for each head into its\nquery-key space separate.\n\nWe could use the fact that both of those two leading axes are basically just batch dimensions, and the existing single matrix multiplication will still work, with one tiny tweak: the current transpose is this:\n\n```\nK_transpose = jnp.transpose(K, axes=(0, 2, 1))\nomega = Q @ K_transpose\n```\n\n...to swap around the last two axes of a three-axis array. With one extra batch dimension, we'll need to take account of that and do this instead:\n\n```\nK_transpose = jnp.transpose(K, axes=(0, 1, 3, 2))\nomega = Q @ K_transpose\n```\n\nThat will be a multiplication of `Q`\n\n, shaped `(batch_size, n_heads, seq_len, d_head)`\n\n,\nwith `K_transpose`\n\n, shaped `(batch_size, n_heads, d_head, seq_len)`\n\n, which\ngives us an `omega`\n\nof the right shape, `(batch_size, n_heads, seq_len, seq_len)`\n\n.\n\nSo, if we can start treating the heads as just another batch dimension, things seem simpler, at least for the attention score calculation. Let's continue down through the single-head code, and then come back later to how we might get the inputs into that double-batched shape.\n\nThe next line after the `omega`\n\ncalculation just scales the attention scores by\na scalar:\n\n```\nomega /= jnp.sqrt(self.d_emb)\n```\n\nThat looked fine, just a broadcast division-by-float. We'd need to change that\n`self.d_emb`\n\nto be `d_head`\n\nin some manner, but that's all.\n\nNext:\n\n```\ncausal_mask = jnp.ones_like(omega, dtype=bool)\n```\n\nThe `jnp.ones_like`\n\nwill give us an array that's `(batch_size, n_heads, seq_len, seq_len)`\n\nfull of `True`\n\ns. That seems reasonable.\n\nThe next step:\n\n```\ncausal_mask = jnp.tril(causal_mask)\n```\n\nWhat will that do? Well, per [the tril documentation](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.tril.html#jax.numpy.tril):\n\nWhen\n\n`m.ndim > 2`\n\n,`jnp.tril`\n\noperates batch-wise on the trailing axes.\n\n...which sounded good. `batch_size`\n\nand `n_heads`\n\nwould be treated as batch axes,\nwhich meant that the next line:\n\n```\ncausal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n```\n\n...would work. Likewise, with the next line:\n\n```\nattention_weights = jax.nn.softmax(causal_omega, axis=-1)\n```\n\n...the axis to apply `softmax`\n\nto is explicitly stated as the last one, which is\nwhat we wanted.\n\nSo at the end of all of those steps, we'd have `attention_weights`\n\nshaped\n`(batch_size, n_heads, seq_len, seq_len)`\n\n, where the last axis had been\nsoftmaxed (softmaxxed?).\n\nThe next line looked a little trickier:\n\n```\nreturn attention_weights @ V\n```\n\nIn the single-head version we had `attention_weights`\n\nof shape `(batch_size, seq_len, seq_len)`\n\n,\nand V of shape `(batch_size, seq_len, d_emb)`\n\n, so multiplying them gives us `(batch_size, seq_len, d_emb)`\n\nIn the new MHA code so far, we had our `attention_weights`\n\nshaped `(batch_size, n_heads, seq_len, seq_len)`\n\n.\nSo in order for the matrix multiplication to work, we'd need `V`\n\nto be shaped\n`(batch_size, n_heads, seq_len, d_head)`\n\n. That would give us a result shaped\nas `(batch_size, n_heads, seq_len, d_head)`\n\n.\n\nAnd conveniently, we'd already decided that the correct shape for `Q`\n\nand for `K`\n\nwas `(batch_size, n_heads, seq_len, d_head)`\n\n. If we could use the same \"magic\" to\ndo the projection into value space -- that is, to get `V`\n\nsuch that the heads formed\na new batch-like axis like we had for `Q`\n\nand `K`\n\n-- then we'd be all set.\n\nSo, at that point, I'd worked out the core of MHA. If we could get all of the\ninputs into the shape `(batch_size, n_heads, seq_len, d_head)`\n\n, and somehow handle\nan output of the shape `(batch_size, n_heads, seq_len, d_head)`\n\n, then we could use\nMHA code something like this:\n\n```\n        # Q and K are (batch_size, n_heads, len_sequence, d_head)\n        # We need to convert K to (batch_size, n_heads, d_head, len_sequence)\n        # and then we get omega (batch_size, n_heads, len_sequence, len_sequence)\n        omega = Q @ jnp.transpose(K, axes=(0, 1, 3, 2))\n\n        omega /= jnp.sqrt(self.d_head)\n\n        causal_mask = jnp.ones_like(omega, dtype=bool)\n        # tril treats all but the last two axes as batches so we're OK here.\n        causal_mask = jnp.tril(causal_mask)\n\n        causal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n\n        # last axis is still OK.\n        attention_weights = jax.nn.softmax(causal_omega, axis=-1)\n\n        # attention_weights is (batch_size, n_heads, len_sequence, len_sequence)\n        # V is (batch_size, n_heads, len_sequence, d_head)\n        # So this will come out as (batch_size, n_heads, len_sequence, d_head)\n        weighted = attention_weights @ V\n```\n\nThe next question was, how do we get our inputs into that shape? We could\nrun them all through separate per-head weights -- that is, have an array with one\nper head, like `W_q[0]`\n\n, `W_k[0]`\n\nand `W_v[0]`\n\nfor the first one. But that, again,\nfelt like it would be failing to take advantage of the GPU properly.\n\nThe solution was to think of how matrix multiplications work. If you multiply two matrices, , the value in the result, in row , and column , is the dot product of row in and column in .\n\nSo, imagine if you wanted to multiply by different versions of , let's call them , , and so on up to . If you imagine a new matrix, , which is basically all the s stacked side-by-side, then the dot-product understanding of multiplication makes it pretty clear that if you did , you would get the results of all of those separate multiplications, also stacked side-by-side. I'll call that kind of matrix a \"striped\" one, for want of a better word.\n\nNow, when we project our inputs into the embedding spaces used for attention, we have code like this:\n\n```\nQ = self.W_q(xs)\n```\n\nWe've initialised the weights, `W_k`\n\nin this case, as an `nnx.Linear`\n\n, so what is\nhappening under the hood here is basically:\n\nThat is, it is just a matrix multiplication. [2](https://www.gilesthomas.com/feed/rss.xml#fn-2)\n\nSo if we imagine that `W_q`\n\nis one of those \"striped\" matrices, holding all of the\nseparate matrices to do the projections for all of the heads in a single one\nshaped `(d_emb, n_heads * d_head)`\n\n, then we could stick with the current code --\nthe\n\n```\nQ = self.W_q(xs)\n```\n\nOur input `xs`\n\nwould be shaped `(batch_size, seq_len, d_emb)`\n\n, so the result would be\n`(batch_size, seq_len, n_heads * d_head)`\n\n, and would have the projections for each head in\nthe same vertical stripes as the separate heads' projection weights.\n\nNow, like PyTorch, JAX allows you to [reshape](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.reshape.html) arrays.\nYou can take one axis of length (say) , and split it into two of\nlengths and respectively -- or, conversely, you can combine two axes of\nlength and to one of .\n\nIf our data had the shape `(batch_size, seq_len, n_heads * d_head)`\n\n, we could reshape it like this:\n\n```\nQ = self.W_q(xs).reshape((batch_size, seq_len, n_heads, d_head))\n```\n\n...and that would split things up. So we'd have Q shaped as `(batch_size, seq_len, n_heads, d_head)`\n\n.\n\nThat's almost what we wanted! We needed `(batch_size, n_heads, seq_len, d_head)`\n\n,\nand a simple transpose could sort that out:\n\n```\nQ = jnp.transpose(\n    self.W_q(xs).reshape(\n        (batch_size, seq_len, n_heads, self.d_head)\n    ),\n    (0, 2, 1, 3)\n)\n```\n\nLikewise for `K`\n\nand `V`\n\n, and that was our inputs sorted.\n\nMoving on to the output; it came from this:\n\n```\nweighted = attention_weights @ V\n```\n\n...and as we worked out above, it was shaped `(batch_size, n_heads, seq_len, d_head)`\n\n.\n\nI remembered that we wanted to run that through a single linear layer to combine\nall of the different heads' outputs into one. It felt like the best way to do that\nwould be to get it back into a \"striped\" layout: `(batch_size, seq_len, n_heads * d_head)`\n\n.\nThis would be something like the inverse of the input-wrangling.\n\nThat would\nneed a reshape, but before I could do that, I'd need to get the axes that needed to be\nmerged next to each other. If the input to the linear layer was going to be\n`(batch_size, seq_len, n_heads * d_head)`\n\n, we'd need to convert it\nfrom `(batch_size, n_heads, seq_len, d_head)`\n\nto `(batch_size, seq_len, n_heads, d_head)`\n\nfirst:\n\n```\njnp.transpose(weighted, (0, 2, 1, 3))\n```\n\n... and then we could just reshape it to `batch_size, len_sequence, n_heads * d_head`\n\n:\n\n```\nstriped_output = jnp.transpose(\n    weighted,\n    (0, 2, 1, 3)\n).reshape(\n    batch_size, len_sequence, self.n_heads * self.d_head\n)\n```\n\nFinally, we could run it through a linear layer, with `in_features`\n\nset to `n_heads * d_head`\n\n,\nand `out_features`\n\nset to `d_emb`\n\n.\n\nI put that all together, and decided to throw something extra into the mix. I remembered that\nRaschka's code had various checks to make sure that `d_head * n_heads == d_emb`\n\n, which\nseemed a little artificial -- I'd read that this was true of GPT-2, but wasn't a necessary\nrestriction for GPT-style models, which makes sense. There's no obvious reason per se why\nthe heads' embedding dimensions should sum up to the higher-level embedding dimensions.\n\nSo I decided initially to just pass in\n`d_head`\n\nand `n_heads`\n\nto the constructor. In my training script I could force them to match\nthe GPT-2 model, but if I wanted to use the code later for something different, I could vary\nthem.\n\nThen I remembered that although the dimensionality of the embedding spaces for the\nquery and the key vectors have to match (because otherwise you can't multiply them\nto work out attention scores with ), the value vector's dimensionality\ncan in theory be different. So I decided to break `d_head`\n\ninto two separate `d_qk`\n\nand\n`d_v`\n\nparameters.\n\n``` python\nclass MultiHeadAttention(nnx.Module):\n\n    def __init__(self, d_emb, n_heads, d_qk, d_v, qkv_bias, rngs):\n        self.n_heads = n_heads\n        self.d_qk = d_qk\n        self.d_v = d_v\n\n        self.W_q = nnx.Linear(d_emb, self.d_qk * n_heads, use_bias=qkv_bias, rngs=rngs)\n        self.W_k = nnx.Linear(d_emb, self.d_qk * n_heads, use_bias=qkv_bias, rngs=rngs)\n        self.W_v = nnx.Linear(d_emb, self.d_v * n_heads, use_bias=qkv_bias, rngs=rngs)\n\n        self.output_projection = nnx.Linear(self.d_v * n_heads, d_emb, use_bias=False, rngs=rngs)\n\n    def __call__(self, xs):\n        batch_size, len_sequence, d_emb = xs.shape\n\n        # For each of the below:\n        # * The initial linear layer projects them to\n        #   (batch_size, len_sequence, d_X * n_heads)\n        #   where X is qk or v as appropriate.\n        # * The reshape makes them (batch_size, len_sequence, n_heads, d_X)\n        # * The transpose makes them (batch_size, n_heads, len_sequence, d_X)\n        Q = jnp.transpose(\n            self.W_q(xs).reshape(\n                (batch_size, len_sequence, self.n_heads, self.d_qk)\n            ),\n            (0, 2, 1, 3)\n        )\n        K = jnp.transpose(\n            self.W_k(xs).reshape(\n                (batch_size, len_sequence, self.n_heads, self.d_qk)\n            ),\n            (0, 2, 1, 3)\n        )\n        V = jnp.transpose(\n            self.W_v(xs).reshape(\n                (batch_size, len_sequence, self.n_heads, self.d_v)\n            ),\n            (0, 2, 1, 3)\n        )\n\n        # Q and K are (batch_size, n_heads, len_sequence, d_qk) per above\n        # We need to convert K to (batch_size, n_heads, d_qk, len_sequence)\n        # and then we get omega (batch_size, n_heads, len_sequence, len_sequence)\n        omega = Q @ jnp.transpose(K, axes=(0, 1, 3, 2))\n\n        omega /= jnp.sqrt(self.d_qk)\n\n        causal_mask = jnp.ones_like(omega, dtype=bool)\n        # tril treats all but the last two axes as batches so we're OK here.\n        causal_mask = jnp.tril(causal_mask)\n\n        causal_omega = jnp.where(causal_mask, omega, -jnp.inf)\n\n        # last axis is still OK.\n        attention_weights = jax.nn.softmax(causal_omega, axis=-1)\n\n        # attention_weights is (batch_size, n_heads, len_sequence, len_sequence)\n        # V is (batch_size, n_heads, len_sequence, d_v)\n        # So this will come out as (batch_size, n_heads, len_sequence, d_v)\n        weighted = attention_weights @ V\n\n        # Transpose to (batch_size, len_sequence, n_heads, d_v),\n        # then reshape to (batch_size, len_sequence, n_heads * d_v)\n        striped_output = jnp.transpose(\n            weighted,\n            (0, 2, 1, 3)\n        ).reshape(\n            batch_size, len_sequence, self.n_heads * self.d_v\n        )\n\n        # Final linear layer to combine\n        return self.output_projection(striped_output)\n```\n\nUnusually for a case where I went off the reservation like this, the whole thing with the embedding space dimensionality didn't cause any problems at all! But there was one small bug in this code, which I didn't discover until later -- we'll come to it by the end of the post.\n\nAt this point, I did another of my short training runs, and:\n\n```\n2026-06-23 17:51:32.094308 Tokens seen: 92,209,152\n2026-06-23 17:51:32.094311 Throughput: 85,682 tokens/second\n2026-06-23 17:51:32.094321 Final train loss: 5.358\n2026-06-23 17:51:32.094323 Done\n```\n\n...with a loss chart that looked like this:\n\nThe training loss at the 937th global step was 5.336, only a tiny bit better than\nthe 5.354 with single-head attention. That was quite possibly within the noise.\nEven though (due to the `d_head * n_heads == d_emb`\n\nrestriction I was enforcing in\nmy training script) the `W_q`\n\n, `W_k`\n\n, and `W_v`\n\narrays were the same size, I was\ncreating that `output_projection`\n\n, which would consume randomness and make things\nvary.\n\nIf I were doing a proper scientific experiment to see if a single layer of MHA beat a single layer of single-head attention, I think I would have run both for more steps to see if the difference became more pronounced later.\n\nBut for the purposes of this post, I decided to move on. My checklist now looked like this:\n\nInside the Transformers blocks, we:\n\nAdding that simple neural network -- the FFN -- seemed like a good next step.\n\nThe [feed forward network](https://www.gilesthomas.com/2025/08/llm-from-scratch-17-the-feed-forward-network)\nis simple; you take the output of the MHA block, run it through a biased\nlinear layer to expand it from `d_emb`\n\nto `4 * d_emb`\n\n, then run it through the\nGELU activation function, then shrink it back down to `d_emb`\n\nwith another linear\nlayer. I didn't really see any value in writing my own implementation of GELU,\ngiven that even in the book we were just given code for an approximation to type\nin. So, using [ jax.nn.gelu](https://docs.jax.dev/en/latest/_autosummary/jax.nn.gelu.html), I\nwrote this:\n\n``` python\nclass TransformersLayer(nnx.Module):\n\n    def __init__(self, d_emb, n_heads, d_qk, d_v, qkv_bias, rngs):\n        self.attention = MultiHeadAttention(d_emb, n_heads, d_qk, d_v, qkv_bias, rngs)\n        self.ffn = nnx.Sequential(\n            nnx.Linear(\n                in_features=d_emb,\n                out_features=d_emb * 4,\n                use_bias=True,\n                rngs=rngs\n            ),\n            jax.nn.gelu,\n            nnx.Linear(\n                in_features=d_emb * 4,\n                out_features=d_emb,\n                use_bias=True,\n                rngs=rngs\n            ),\n        )\n\n    def __call__(self, xs):\n        shortcut = xs\n        att = self.attention(xs)\n        post_attention = shortcut + att\n        fed_forward = self.ffn(post_attention)\n        return fed_forward + post_attention\n```\n\nNote that I added in a shortcut connection around the FFN as well, so that it didn't overwrite what was there, but only \"added on its notes\".\n\nI kicked that off, and it ran for ten minutes or so, but then OOMed:\n\n```\n2026-06-23 18:20:01.376377 Saving checkpoint\n 51%|██████████████████████████████████████████████████████▊                                                    | 481/938 [10:20<11:47,  1.55s/it, loss=5.758, tps=76,185]W0623 18:20:12.602631 2860192 bfc_allocator.cc:514] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.93GiB (rounded to 3149744640)requested by op\nIf the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.\n```\n\nAdding `TF_GPU_ALLOCATOR=cuda_malloc_async`\n\ndidn't help. I spent some time trying\nto dig into what might be causing it, but eventually noticed something interesting:\nin `nvtop`\n\n, the VRAM usage was consistently 75% throughout.\n\nNow I knew that JAX pre-allocates 75% of VRAM when it starts up, but I'd been assuming that it would try to grab more if it needed it. It turned out I was wrong with that assumption -- it grabs 75%, but that's all you ever get!\n\nThe solution turned out to be the [ XLA_PYTHON_CLIENT_MEM_FRACTION](https://docs.jax.dev/en/latest/gpu_memory_allocation.html) environment variable.\nIf you set that to, say,\n\n`0.90`\n\n, then JAX will pre-allocate 90% of the VRAM, and\nyou can use all of that. (You can also make it allocate as-needed with\n`XLA_PYTHON_CLIENT_PREALLOCATE=false`\n\n, and there are various other settings you\ncan control with other environment variables on that linked page).Anyway, setting it to `0.90`\n\nto grab 90% of VRAM worked, and I was able to get\na successful run:\n\n```\n2026-06-24 00:29:34.864880 Tokens seen: 92,209,152\n2026-06-24 00:29:34.864882 Throughput: 77,596 tokens/second\n2026-06-24 00:29:34.864900 Final train loss: 5.341\n2026-06-24 00:29:34.864902 Done\n```\n\nThe loss chart was this:\n\n...and the training loss at global step 937 was 5.295, compared to the 5.336 from MHA alone. Another tiny improvement, another one that could have been in the noise. Again, if I were doing a proper experiment, I'd do a longer run, but for now, I decided to move on.\n\nThe checklist looked like this:\n\nInside the Transformers blocks, we:\n\nNow, my gut instinct was that the layer normalisation inside the Transformers blocks was of most value as a way of stabilising training over deep networks. And with one layer, it didn't seem like the right time to add it. Instead, I decided to add on multiple layers.\n\nFor GPT-2 small, you have 12 layers. That was already being passed in to my `GPTModel`\n\n's\n`__init__`\n\nmethod as `n_layers`\n\n, so I just replaced this:\n\n```\nself.transformers_layer = TransformersLayer(\n    d_emb, n_heads, d_qk, d_v, qkv_bias, rngs\n)\n```\n\n...with this:\n\n```\nself.transformers_layers = nnx.Sequential(\n    *(\n        TransformersLayer(\n            d_emb, n_heads, d_qk, d_v, qkv_bias, rngs\n        )\n        for _ in range(n_layers)\n    )\n)\n```\n\n...and then just renamed it where it was called; this:\n\n```\ntransformed = self.transformers_layer(input_embeddings)\n```\n\n...became this:\n\n```\ntransformed = self.transformers_layers(input_embeddings)\n```\n\nI kicked it off, and it completed! However, the loss chart was telling:\n\nOuch. Loss started dropping quite nicely, but then things got out of control and it settled down at a loss that was essentially that of a random model. At step 937, we were at 10.75, so just a hair less than the 10.82 that randomly guessing next tokens would give.\n\nWell, LayerNorm is specifically meant to stabilise training, and the checklist looked like this:\n\nInside the Transformers blocks, we:\n\n...and the only remaining step was that LayerNorm in the Transformers blocks, so it was time to add it in!\n\nAs per the checklist, we do the LayerNorm after we've taken our copy for the shortcut connection, just before MHA, and then likewise after the second shortcut copy, before the FFN. As I understand it, this was a GPT-2 innovation -- previously, people had done normalisation after those steps, but this pre-norm setup turned out to work better.\n\nThe code changes were simple. I added two `LayerNorm`\n\nmodules to the\n`TransformersLayer`\n\nclass, and then called them in the appropriate places (taking\nthe opportunity to tidy up the variable naming in the forward pass while I was there):\n\n``` python\nclass TransformersLayer(nnx.Module):\n\n    def __init__(self, d_emb, n_heads, d_qk, d_v, qkv_bias, rngs):\n        self.attention_norm = LayerNorm(d_emb)\n        self.attention = MultiHeadAttention(d_emb, n_heads, d_qk, d_v, qkv_bias, rngs)\n\n        self.ffn_norm = LayerNorm(d_emb)\n        self.ffn = nnx.Sequential(\n            ...\n        )\n\n    def __call__(self, xs):\n        shortcut = xs\n        xs = self.attention_norm(xs)\n        xs = self.attention(xs)\n        xs = xs + shortcut\n\n        shortcut = xs\n        xs = self.ffn_norm(xs)\n        xs = self.ffn(xs)\n        return xs + shortcut\n```\n\nI kicked it off and ran it, and got these results:\n\n```\n2026-06-24 03:07:51.128966 Tokens seen: 92,209,152\n2026-06-24 03:07:51.128969 Throughput: 23,399 tokens/second\n2026-06-24 03:07:51.128979 Final train loss: 5.359\n2026-06-24 03:07:51.128981 Done\n```\n\nThat certainly looked much healthier!\n\nHowever, when I looked at the loss at step 937, it was 5.311 -- a tiny bit higher than the single-layer MHA example, which got 5.295.\n\nI'd been willing to play a bit fast and loose with this loss number and allow myself\nto accept a win when the loss went down a tiny bit, even if it was such a small\namount that it could have been within the noise. But *increasing* loss -- even if\nit could *also* be within the noise -- was a step too far.\n\nI decided that in this specific case, I'd be strict and test the hypothesis that longer training runs would demonstrate an improvement between one single layer without pre-norm, and multiple layers with pre-norm.\n\nI had to remember that these training runs would not be comparable with the earlier ones.\nIn the training script, I [had a learning rate schedule like this](https://www.gilesthomas.com/2026/06/llm-from-scratch-34a-building-a-jax-training-loop-for-an-llm-training-run#learning-rate-scheduling):\n\nThat straight-line warmup period and the following cosine decay were 5% and 95% of the training run respectively, which meant that (for example) global step 937 of the short runs we had been doing would be at a completely different point in the schedule than the same step would in these longer runs.\n\nHowever, they would be comparable to each other, and that was what mattered.\n\nAfter some humming and hawing, I decided that a full Chinchilla-optimal (for the full model) training run over 3,260,190,720 tokens, rounded up to fit into a round number of global steps, would be a nice experiment. I expected it to run comfortably overnight for the single-layer run, and take a bit less than two days for the multi-layer one. So I kicked off the first.\n\nJust over 11 hours later:\n\n```\n2026-06-24 16:39:52.178045 Tokens seen: 3,260,252,160\n2026-06-24 16:39:52.178049 Throughput: 81,733 tokens/second\n2026-06-24 16:39:52.178058 Final train loss: 4.324\n2026-06-24 16:39:52.178061 Done\n```\n\nHere's the loss chart:\n\nThe last checkpointing period in that run ended at global step 33,164, and the training loss then was 4.165 -- indeed, it had been at around 4.17 for quite some time, though the trend still seemed to be a tiny bit downward.\n\nSo then I kicked off a run of the full version -- multiple layers, with pre-norm in the Transformers blocks. Just over 37 hours later:\n\n```\n2026-06-26 06:55:39.956609 Tokens seen: 3,260,252,160\n2026-06-26 06:55:39.956614 Throughput: 24,151 tokens/second\n2026-06-26 06:55:39.956625 Final train loss: 3.637\n2026-06-26 06:55:39.956629 Done\n```\n\nThe \"Final train loss\" line at the end said it all, really! But here's the loss chart:\n\n...and the loss at step 33,164 was 3.399. Definitely quite an improvement over the 4.165 that a single layer got.\n\nAgain, at some point I might do the equivalent tests for the earlier results where improvements appear to be pretty much in the noise. It would be good to be sure that the changes really did have the impact I think they did.\n\nBut for now: our checklist was looking like this:\n\nInside the Transformers blocks, we:\n\nEverything was checked off. So was this journey over?\n\nWell, there was one thing that the original PyTorch code had that my new code didn't: dropout.\n\nI'd found in my lengthy [interventions experiments](https://www.gilesthomas.com/2026/02/llm-from-scratch-32a-interventions-baseline-model) that\ndropout seemed to make models worse. It was, I felt, a smart idea back in the days\nwhen people had little data and did multiple epochs, each sweeping over everything,\nbut it made less sense nowadays with single-epoch training runs over very large datasets.\n(Though I do have [some intuitive ideas about why it could still help](https://www.gilesthomas.com/2025/03/dropout-and-mandatory-vacation).)\n\nStill, it would be good to show that it harmed loss for this model as well.\n\nChecking my notes, I found that there were four places where dropout was applied:\n\nThe changes are tiny and rather dotted around the\ncode, so rather than showing you isolated bits of code, if you'd like to see it you\ncan take a look at [the code at this point](https://github.com/gpjt/jax-gpt2-from-scratch/blob/5fe6e745b55a33a3b55d3d7ce39f480e29ed0d77/gpt.py)\nand search for \"dropout\".\n\nWhen I started running that, I got an error when saving the first checkpoint:\n\n```\nTypeError: JAX array with PRNGKey dtype cannot be converted to a NumPy array. Use jax.random.key_data(arr) if you wish to extract the underlying integer array.\n```\n\nThis was happening deep inside the bowels of Safetensors, but it made a lot of sense.\nThe `nnx.Dropout`\n\nobject needs to keep track of the state of the random number generator,\nand that meant that the `to_flat_state`\n\nfunction that [I was using](https://www.gilesthomas.com/2026/06/flax-and-safetensors)\nmight return a structure that had something that contained that state, and was not\ncompatible with Safetensors.\n\nI decided that I'd cheat a little bit here. If I skipped the dropout layers when I saved my checkpoints, like this:\n\n```\n    for tuple_key, array in flat_state:\n        key = \".\".join(str(key) for key in tuple_key)\n        if \"dropout\" not in key:\n            simple_dict[key] = array\n```\n\n...then I'd be able to save them. This would have a problem -- if I restarted from a checkpoint, the dropout pattern after the restart would mirror the dropout pattern from the start of the training run, because the random seed it started with would not have come from the checkpoint, but just the initialisation code.\n\nI felt that this would not have a serious impact, though, and given that I'd not had to restart from checkpoints so far, I (wrongly, as it turned out) decided it wouldn't matter.\n\nI kicked off the run, and... after four hours, it OOMed. I cursed, decided that I'd nurse this run through anyway (despite my dropout checkpointing concerns), and kicked it off again.\n\nThree hours later, it OOMed again. I happened to be away from home at the time, logging in to\nmy machine remotely (thanks, [Tailscale](https://tailscale.com/)!), and\non looking at `nvtop`\n\n, I realised that the X window system on my machine was using a\ngig or so of VRAM.\n\nI was running the training run in a `tmux`\n\nsession, which meant\nthat I could kill X and not lose state, so I did that, and adjusted the\n`XLA_PYTHON_CLIENT_MEM_FRACTION`\n\nenvironment variable I was using -- it had been\n0.90, so I bumped it up to 0.95. I kicked it off again, and...\n\n```\n2026-06-28 22:06:47.669676 Tokens seen: 2,640,052,224\n2026-06-28 22:06:47.669683 Throughput: 23,019 tokens/second\n2026-06-28 22:06:47.669691 Final train loss: 3.776\n2026-06-28 22:06:47.669694 Done\n```\n\nNote that the tokens seen only relates to the period since the restart, which is why it was lower.\n\nOne more loss chart:\n\n...and the training loss at step 33,164 was 3.524, higher enough than the 3.399 I got without dropout that I was comfortable that it wasn't in the noise. That was very reassuring.\n\nOnce again, if this was a proper scientific experiment I'd fix the issue with saving dropout, and run it completely from scratch -- or, at least, run it all the way through from scratch without restarts, even if I had to try several times to get it done.\n\nBut I don't think that \"replaying\" dropout would make the loss any worse. And for this experiment, I felt this was enough.\n\nSo: checklist complete. GPT-2 model coded up. It was time for some evals!\n\nI wanted to evaluate these models against the ones I got using the old PyTorch code:\nspecifically, the [last local training run](https://www.gilesthomas.com/2026/04/llm-from-scratch-32k-interventions-training-our-best-model-locally-gradient-accumulation)\nthat used exactly the same training hyperparameters, and only differed in that\nit was trained using AMP -- 32-bit floats in general, but using 16-bit where the\nframework thought it would not be harmful.\n\nIn order to do *exactly* the same evals, I decided it would be easiest to write a\nconversion script to take the Safetensors files written to my JAX checkpoints, and write\nout new files that were compatible with the PyTorch model code -- then I'd be able to\nuse the original PyTorch eval code. I\n[put something together](https://github.com/gpjt/jax-gpt2-from-scratch/blob/fd170272e54844229642619b9c69993f0f18778f/convert_model_to_pytorch.py),\nconverted my last two models -- the full runs with and without dropout -- and tried to\nload them up.\n\nUnfortunately there was an error:\n\n```\nRuntimeError: Error(s) in loading state_dict for GPTModel:\n    Missing key(s) in state_dict: \"trf_blocks.0.att.out_proj.bias\", \"trf_blocks.1.att.out_proj.bias\", \"trf_blocks.10.att.out_proj.bias\", \"trf_blocks.11.att.out_proj.bias\", \"trf_blocks.2.att.out_proj.bias\", \"trf_blocks.3.att.out_proj.bias\", \"trf_blocks.4.att.out_proj.bias\", \"trf_blocks.5.att.out_proj.bias\", \"trf_blocks.6.att.out_proj.bias\", \"trf_blocks.7.att.out_proj.bias\", \"trf_blocks.8.att.out_proj.bias\", \"trf_blocks.9.att.out_proj.bias\"\n```\n\nYou might remember that back when I went through multi-head attention, I mentioned\nthat I'd made a mistake. Somehow, I'd misremembered, and thought that the output projection -- the one\nthat mixes together all of the different heads' outputs -- was a linear layer without\nbias, despite [my original notes](https://www.gilesthomas.com/2025/04/llm-from-scratch-12-multi-head-attention) being\nperfectly clear that it *did* have bias.\n\nThe good news was that if I disabled bias in the PyTorch code, I could load the safetensors files that I had. So the two models I'd trained so far were not useless, and could actually work as a kind of natural experiment into the benefits of having that bias there.\n\nBut anyway, in order to do things properly, I was going to need to fix the bug and train yet another model.\n\nThe fix was simple, I just replaced this (in `MultiHeadAttention`\n\n):\n\n```\nself.output_projection = nnx.Linear(self.d_v * n_heads, d_emb, use_bias=False, rngs=rngs)\n```\n\n...with this:\n\n```\nself.output_projection = nnx.Linear(self.d_v * n_heads, d_emb, use_bias=True, rngs=rngs)\n```\n\nThen it was time to kick off yet another training run. After another 37 hours:\n\n```\n2026-07-05 10:09:22.147819 Tokens seen: 3,260,252,160\n2026-07-05 10:09:22.147823 Throughput: 24,072 tokens/second\n2026-07-05 10:09:22.147832 Final train loss: 3.650\n2026-07-05 10:09:22.147834 Done\n```\n\n...with this loss chart:\n\n...and the training loss at step 33,164 was 3.398 -- almost exactly the same as the 3.399 that I got in the no-dropout training run without MHA bias above!\n\nWell, now it really was time for the evals.\n\nI updated [my conversion script](https://github.com/gpjt/jax-gpt2-from-scratch/blob/9ab47bc589973b3d8d18a5248b405ef3f17f09fe/convert_model_to_pytorch.py)\nto handle the bias on the MHA output projections, and used it to convert the three\nmodels -- the un-biased ones, with and without dropout, and the biased one, without -- to\nthe PyTorch format, then ran the loss test that I had been using to compare the old models\non each.\n\nHere are the results, compared to the previous models, and OpenAI's:\n\n| Test loss | |\n|---|---|\n| OpenAI weights: medium | 3.231442 |\nJAX, with MHA bias, no dropout |\n3.418784 |\nJAX, no MHA bias, no dropout |\n3.420089 |\nJAX, no MHA bias, with dropout |\n3.476802 |\n| OpenAI weights: small | 3.499677 |\n`1xrtx3090-stacked-interventions` |\n3.538161 |\n`8xa100m40-stacked-interventions-1` |\n3.577761 |\n| Cloud FineWeb, 8x A100 40 GiB | 3.673623 |\n`1xrtx3090-baseline` |\n3.683835 |\n`8xa100m40-baseline` |\n3.691526 |\n| Cloud FineWeb, 8x H100 80 GiB | 3.724507 |\n| Cloud FineWeb, 8x A100 80 GiB | 3.729900 |\n| Cloud FineWeb, 8x B200 160 GiB | 3.771478 |\n| Local FineWeb train | 3.943522 |\n| Local FineWeb-Edu extended train | 4.134991 |\n| Local FineWeb-Edu train | 4.166892 |\n\nThat was a pretty amazing result -- I'd clearly proven that JAX trains much better models than PyTorch! 3.5% better in the best case.\n\nWell, OK, no.\n\nMy guess is that the difference was probably something\nlike better luck with the initial weights on the JAX side, plus\n[the improvement from not using AMP](https://www.gilesthomas.com/2026/04/llm-from-scratch-32h-interventions-full-fat-float32).\nAnyway, the important thing was that the JAX models were in the same kind of loss\nrange as the PyTorch ones -- and while a 3.5% improvement in loss was more variation\nthan I'd been expecting, it was definitely the right ballpark.\n\nNow, one thing I had found in the past was that the OpenAI weights -- and some of\nmy own models, like the Fineweb-Edu ones -- were\n[consistently better](https://www.gilesthomas.com/2026/04/llm-from-scratch-32l-interventions-instruction-fine-tuning-tests)\nat an instruction fine-tuning test than their test loss scores would indicate. Would that hold here?\n\nThe IFT eval code fine-tuned\neach model on the [Alpaca](https://crfm.stanford.edu/2023/03/13/alpaca.html) dataset until\nvalidation loss started rising, then used the model prior to the start of the rise to\ngenerate responses for a test set. These were saved, and then run past an OpenAI model\nso that they could be compared with each other:\n\n```\nYou are judging the comparative capabilities of a number of different LLM\nmodels.  They have been trained to follow instructions.\n\nThe input was this:\n\n`\n{input}\n`\n\nAn example correct output is this:\n\n`\n{correct_output}\n`\n\nPlease produce a score of between 0 and 100 for each model, and respond\nwith a JSON structure like this (note that the number of models may differ\nfrom this example):\n\n`\n{\n    \"Model 1\": {\"score\": XXX, \"comments\": \"optional comments\"},\n    \"Model 2\": {\"score\": YYY, \"comments\": \"optional comments\"},\n    \"Model 3\": {\"score\": ZZZ, \"comments\": \"optional comments\"}\n}\n`\n\n...where the XXX, YYY and ZZZ are the scores for the respective models.\nYou can optionally add the \"comments\" field if you want to explain your\nreasoning.\n\nHere are the models' responses:\n\n# Model 1\n\n{model 1 response}\n\n# Model 2\n\n{model 2 response}\n\n# Model 3\n\n{model 3 response}\n```\n\n...with the model order randomly changed for each query to avoid any position bias.\n\nThe methodology seemed solid, but I was uncertain about the \"train until loss starts rising\", as it meant that different models had wildly different amounts of fine-tuning -- between two and seven epochs.\n\nOn the one hand it felt \"unfair\" to certain models that they'd get less training than others. On the other hand, if the less-trained models had been trained past the point where their validation loss started rising, then assuming that loss would continue to rise, further training would actually be a disadvantage rather than an advantage.\n\nI decided to stick with the original plan, and train until validation loss started rising. I did, however, switch the judge model from the GPT 5.4 that I used in my last IFT test to GPT 5.5. Here are the results:\n\n| Test loss | IFT epochs | IFT score | IFT rank | |\n|---|---|---|---|---|\n| OpenAI weights: medium | 3.231442 | 2 | 41.62 | 1 |\nJAX, with MHA bias, no dropout |\n3.418784 | 4 | 19.25 | 4 |\nJAX, no MHA bias, no dropout |\n3.420089 | 3 | 14.66 | 11 |\nJAX, no MHA bias, with dropout |\n3.476802 | 4 | 12.94 | 15 |\n| OpenAI weights: small | 3.499677 | 2 | 26.73 | 2 |\n`1xrtx3090-stacked-interventions` |\n3.538161 | 4 | 17.79 | 6 |\n`8xa100m40-stacked-interventions-1` |\n3.577761 | 4 | 10.29 | 16 |\n| Cloud FineWeb, 8x A100 40 GiB | 3.673623 | 7 | 20.71 | 3 |\n`1xrtx3090-baseline` |\n3.683835 | 6 | 15.11 | 9 |\n`8xa100m40-baseline` |\n3.691526 | 4 | 14.74 | 10 |\n| Cloud FineWeb, 8x H100 80 GiB | 3.724507 | 4 | 13.25 | 14 |\n| Cloud FineWeb, 8x A100 80 GiB | 3.729900 | 4 | 14.50 | 12 |\n| Cloud FineWeb, 8x B200 160 GiB | 3.771478 | 4 | 16.03 | 8 |\n| Local FineWeb train | 3.943522 | 7 | 13.73 | 13 |\n| Local FineWeb-Edu extended train | 4.134991 | 7 | 16.70 | 7 |\n| Local FineWeb-Edu train | 4.166892 | 7 | 18.68 | 5 |\n\nMore interesting datapoints! As before, you can see that low loss is not particularly well-correlated with a high score on this instruction fine-tuning test. The OpenAI weights continue to lead the pack, and while one of our new JAX models did quite well, it's still beaten by the Cloud FineWeb, 8x A100 40 GiB model.\n\nBut what was important here, just as with the loss, was that the new JAX models landed in the same ballpark as the PyTorch ones. They did, and so I could be confident that they were doing essentially the same thing.\n\nAnd that meant that, after 18 months, I had reached the end of my LLM from scratch journey.\n\nIt's been [a long trek](https://www.gilesthomas.com/llm-from-scratch).\n\nI [started reading](https://www.gilesthomas.com/2024/12/llm-from-scratch-1) \"Build a Large Language Model (from Scratch)\"\non 22 December 2024. I was planning to breeze through over the Christmas break,\nbut somehow it morphed into being a curriculum onto which I could hang projects\nto learn the fundamentals of LLMs, beyond what was in the book.\n\nIn May 2025, I had my first real conceptual breakthrough when I realised that\n[attention heads are (individually) dumb](https://www.gilesthomas.com/2025/05/llm-from-scratch-13-taking-stock-part-1-attention-heads-are-dumb),\nand as I continued, the second big one came later on in the same month, when the concept of\nembeddings as being [projections between vocab space and embedding space](https://www.gilesthomas.com/2025/05/llm-from-scratch-15-from-context-vectors-to-logits)\n(and the converse projection in the other direction that happens in the LLM's output head)\nbecame clear.\n\nIn August I had the first moment where I felt that the standard teaching approach to\nLLMs might not be the full story; shortcut connections are normally explained as\na way to fix vanishing gradients, while I felt that a better way to see them was a\nway to allow attention and the FFN to \"annotate\" the existing information, similarly\nto how Jewish scholars have [annotated the original text of the Talmud](https://www.gilesthomas.com/2025/08/llm-from-scratch-18-residuals-shortcut-connections-and-the-talmud).\n(The results in this post seem to point in that direction, given how even a single\nlayer of attention was massively helped by adding them.)\n\nBy early December, I had essentially finished the book, and felt I wanted to try\nto [train my first base model from scratch on my RTX 3090](https://www.gilesthomas.com/2025/12/llm-from-scratch-28-training-a-base-model-from-scratch).\nIt worked, and wasn't far off the quality of the original GPT-2 small. I was really surprised\nthat I could do that with consumer hardware, and became interested (perhaps obsessively so)\nwith whether I could match OpenAI's weights.\n\nIn January 2026, I [trained a model using DDP on Lambda Labs](https://www.gilesthomas.com/2026/01/llm-from-scratch-29-ddp-training-a-base-model-in-the-cloud),\nand then spent the following months training model after model, trying to work\nout which interventions -- learning rate scheduling, gradient clipping, etc -- would\nimprove the loss. I wrapped that up in [late April](https://www.gilesthomas.com/2026/04/llm-from-scratch-32m-interventions-conclusion),\nwith the interesting finding that although I'd been able to get the test loss pretty\nlow, that didn't seem to map cleanly to performance in my instruction fine-tuning\ntests. In other words, Loss Number Goes Down is an interesting technical game to play,\nbut doesn't cleanly map to real-world performance.\n\nThe final step was this post, and [the previous one](https://www.gilesthomas.com/2026/06/llm-from-scratch-34a-building-a-jax-training-loop-for-an-llm-training-run)\n-- could I, using my notes, implement GPT-2 completely from scratch in JAX without referencing\nthe book?\n\nAnd as you've read, the answer was a definite yes!\n\nOf course, as with any long-running project, there are some loose ends -- from this\npost alone, there's the interesting fact that JAX trained faster than PyTorch (perhaps\n`torch.compile`\n\ncould close the gap?) and had a larger possible batch size for full-fat 32-bit.\nAnd the fact that fixing the multi-head attention bias bug didn't seem to help with the\nloss much was interesting too.\n\nBut those are really details, and there's so much beyond them to learn. Longer-context LLMs: position embedding improvements like RoPE,\nefficiency tricks like flash attention and attention variants like DSA. Mixture\nof experts models. How do optimisers really work? ([Do they work?](https://x.com/mgostIH/status/2024087619756318866))\nAnd plenty more.\n\nSo it's time to draw a line under this series, and start thinking about what comes next. It's been a blast; if you've been reading along, I hope it's been as useful (and fun) to read as it was to write.\n\nAnd as always, comments, questions and corrections very welcome below.\n\nOn looking back at Raschka's code, after having worked through all of this, there's a slight difference. I do this:\n\n```\nnormalized = (xs - means) / (stds + 1e-5)\n```\n\n...whereas he does this:\n\n```\nnorm_x = (x - mean) / torch.sqrt(var + self.eps)\n```\n\nNow, the standard deviation is the square root of the variance, so if you ignore\nthe small numbers -- `1e-5`\n\nin my case, and `self.eps`\n\nin his -- the calculations\nare the same. But there is a difference once those are taken account of.\nI don't think it's large enough to have any serious effect in these runs, though. [↩](https://www.gilesthomas.com/feed/rss.xml#fnref-1)\n\nIn PyTorch, linear layers are stored as the transpose of the matrix that would allow you to do that, so it would be:\n\nAlso, note that for simplicity (heh) I'm disregarding bias in this discussion. [↩](https://www.gilesthomas.com/feed/rss.xml#fnref-2)", "url": "https://wpnews.pro/news/writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in", "canonical_source": "https://www.gilesthomas.com/2026/07/llm-from-scratch-34b-building-and-training-gpt-2-small-in-jax", "published_at": "2026-07-08 17:59:57.661129+00:00", "updated_at": "2026-07-08 17:59:59.776284+00:00", "lang": "en", "topics": ["large-language-models", "artificial-intelligence", "ai-research", "ai-products"], "entities": ["Giles Thomas", "Sebastian Raschka", "JAX", "PyTorch", "GPT-2", "RTX 3090", "OpenAI", "X"], "alternates": {"html": "https://wpnews.pro/news/writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in", "markdown": "https://wpnews.pro/news/writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in.md", "text": "https://wpnews.pro/news/writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in.txt", "jsonld": "https://wpnews.pro/news/writing-an-llm-from-scratch-part-34b-from-bigrams-to-gpt-2-one-component-at-a-in.jsonld"}}