Giles' blog

Benchmarking Qwen 3.6 35B MoE (3B active) on an RTX 3090

Fri, 24 Jul 2026 23:30:00 +0000

I mentioned I'd got a new RTX 3090 on a group chat, and a friend said:

I know this is not really your thing... but let me know how quickly it runs Qwen 3.6 35bn MoE. With only 24gb of VRAM you’ll need to use a 4-bit quantized version and you won’t get a massive context window. But it should still be pretty cool.

He's right that it's not really been my thing -- I've been focusing on my own LLMs recently. I decided to dig in a little, and in particular to play with Llama.cpp, which I haven't used for a while. And then things got a tad out of control, and I wound up doing some relatively detailed benchmarking.

The headline results: I downloaded Unsloth's UD-IQ4_NL_XL quantisation of the model from Hugging Face. With that, using the default Arch build of Llama.cpp, which uses Vulkan under the hood:

Using the GPU only, I was able to get the model to generate at just over 120 tokens per second, and it was able to process the prompt at just less than 2,800 tok/s. However, having the whole model on the GPU didn't leave that much space for the context window -- it was constrained to about 50,000 tokens, compared to the model's native context length of 262,144.
Offloading the FFNs for the first 12 of the model's 40 layers to the CPU managed to reclaim enough VRAM to be able to get the full context length; however, with that setup, things were -- unsurprisingly -- slower. I got just over 65 tok/s for generation, and 600 tok/s for the prompt.

Compiling Llama.cpp myself, in order to get the full CUDA version, helped a lot:

With everything on the GPU, I got 140 tok/s for generation and over 3,300 tok/s for the prompt. That was with a context window of 89,600. So everything was better :-)
It was also easier to get to the full context length; that needed just 10 layers' FFNs to be offloaded, and at that point I was getting 89 tok/s for generation, and about 1,100 tok/s for the prompt.

That was a pretty impressive improvement.

But it also showed that the lack of VRAM on the 3090 really does hurt. The friend who asked me about this is running an Intel Arc B70 Pro, with 32 GiB VRAM. He can (of course) fit the whole 4-bit quantised model on there without any context window reductions, and he says this about throughput:

It starts off at 75-80. but drops into the high 50s as the context window expands

That's worse than the CUDA-with-offload results above, though I did limit my testing to a 2,457-token prompt, with 6,144 tokens generated. The RTX 5090 has 32 GiB VRAM and fast Nvidia processing -- I imagine things would be a lot better there. Downside: it costs more than three times what you pay for a second-hand 3090 (and double the Arc B70 Pro).

Anyway, in the rest of this post, I'll give more details of the benchmark, including charts showing the performance at different numbers of offloaded layers, and comparisons of the Vulkan vs CUDA versions of Llama.cpp for this model.

One caveat: Qwen 3.6 is a multimodal model. That means that it has an extra component to map graphical inputs to tokens that can be consumed by the LLM. I didn't realise this as I ran these tests, but you can actually tell it not to use that component if you're only working with text inputs -- for example, --no-mmproj will do the trick with Llama.cpp. If you're working with text-only stuff, you might want to play with that to see if you can get better results -- it will free up some VRAM and might allow you to get longer context lengths, or better performance.

But anyway, as with many of my posts, this is a tidied up version of my lab notes, so you can share my learning journey with me -- but if that sounds pointless, and you've landed here because you want to run this model on your own RTX 3090, you can click this link to jump right to the results.

The model

Qwen3.6-35B-A3B is a Mixture of Experts model with 35B total parameters, 3B active. The tensor type is BF16, so it's two bytes per parameter. That's ~70 GiB of space required for the model, just for the weights, with no space for attention matrices and the like. We're going to need a quantised version.

I'll dig deep into MoE models at some point, but essentially they work by having multiple different FFN blocks for each Transformers layer. After attention, the context vectors are routed to a subset of those blocks, depending on their content ¹. The number of active parameters is how many are actually used when running a single token through.

Now, that number, 3B, looked really crazy to me as soon as I saw it. Qwen 3.6, as you can see from the model card linked above, has a vocabulary size of 248,320, and an embedding dimensionality of 2,048. It doesn't use weight-tying, so that means that the embedding layer and the output head come out at a hefty 508,559,360 parameters each.

Both the embedding layer and the output head are required for every token, of course, so that means that of those 3B active parameters, more than 1B are used just getting tokens into and out of the model! Only 2B active parameters are left to handle all of the attention and the active FFNs. That is an imbalance of the same level as you get with GPT-2 small, and was surprising to see.

Well, let's think about VRAM requirements. Naively you might think that you only need to keep the active parameters in VRAM, swapping FFN experts in and out as needed. Unfortunately that doesn't work.

Let's say you feed in "The fat cat sat on the", hoping for a completion. The whole sequence is processed in parallel (that being the point of using a GPT-style LLM rather than, say, an RNN). You'll need whichever experts are required for all six tokens in the prompt for the first layer, and likewise for all of the context vectors for the prompt in each of the later layers.

So, realistically, for a non-trivial prompt, you can't "swap out" the inactive parameters; an MoE buys you lower usage in terms of processing, but not in terms of memory. ² That said, as we'll see later, you can at least pull stuff out of VRAM into normal RAM, and get some advantages that way.

As my friend said, a 4-bit quant would be needed -- perhaps with that, the whole model, both active and inactive parameters, would fit into 35 / 2 = 17.5 GiB of VRAM, leaving 6.5 GiB for activations, attention matrices, and so on? As it turned out, it needed a bit more, and it's worth taking a look into why.

Quantisation

I must admit that I'd never really looked into quantisation prior to playing with this, and had naively assumed that you basically just took (say) a BF16 model, scaled the parameters down so that they were (say) FP4, tweaked the computations a bit and then ran the result.

That was completely wrong! I'll have to dig into it further in the future, but my new working mental model is that it's more like lossy compression. The weights are stored in a "compressed" format -- one that is designed so that they can be quickly and cheaply "decompressed" by code running on the target platform.

So, speaking very loosely, when running an unquantised model, we might have CUDA code doing something like this:

Load BF16 weights from VRAM
Load BF16 data from another bit of VRAM
Do a matrix multiplication of one by the other.
Store the results into VRAM as BF16

...a quantised model might operate more like this:

Load quantised weights from VRAM
"Decompress" them to BF16
Load BF16 data from another bit of VRAM
Do a matrix multiplication of one by the other.
Store the results into VRAM as BF16

Understanding that minimal amount clarified three things for me:

Why people talk about things like "4.1-bit quants". We're talking about the average number of bits per weight in the "compressed" model.
Why there are different quants of about the same size for a given model; for example, when we get on to looking at quantised versions of this model later, you'll see 4-bit ones called UD-IQ4_NL, UD-Q4_K_M, UD-IQ4_NL_XL, and UD-Q4_K_XL, with different sizes in terms of GiB. Each of those is using a different set of trade-offs in the quantisation process, so will perform differently on different tasks. I gather that picking the right one for any given task is a bit of an art form. Hugging Face have a summary of the different types they support here, which has some explanation of the differences, but decoding what they're saying there is beyond what I've learnt at this point.
How GPUs can run quants with bits-per-parameter levels that they don't support for calculations. For example, the RTX 3090 supports 32-bit and two forms of 16-bit (float16 and BF16) plus integer operations of various bittednesses ³. 4-bit formats like FP4 are only supported in more recent cards like the RTX 5090. But, of course, we're not doing any 4-bit computations -- it's all in some format that the GPU can handle natively.

So, with that (minimal) level of understanding, it was time to dig in!

Getting Llama.cpp up and running

There are a bunch of different ways that people run LLMs locally, so I needed to work out which one would be likely to give the best results. One thing I noticed when googling around for things like "Qwen3.6-35B-A3B RTX 3090" was that pretty much everyone appeared to be using Llama.cpp. That wasn't something I'd used for a couple of years, and I've repaved my machine multiple times since, so it was time to install it.

I run Arch Linux, and there's an OS package for it, so I installed it using the instructions there -- specifically, the ones to run inference using CUDA:

sudo pacman -S llama-cpp ggml cuda

Now it was time to try a model. I decided to start off with a small one that would fit onto my GPU without any quantisation, just to work out any bugs. I needed the model to be in GGUF format (that being what Llama.cpp uses). GGUF stores tensors and metadata -- the actual architectures of the LLMs Llama.cpp supports are baked in to the tool's source -- but, of course, it supports Qwen3, and so a nice simple non-MoE small model that looked like it would fit onto my GPU was Qwen/Qwen3-4B-GGUF.

The help page for that model gave this example command to run it:

./llama-cli -hf Qwen/Qwen3-4B-GGUF:Q8_0 --jinja --color -ngl 99 -fa -sm row \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5 \
    -c 40960 -n 32768 --no-context-shift

That was out-of-date and gave various errors (which was fair enough, the model being from late December 2025 -- all of seven months old!). The fixes were simple enough:

llama-cli -hf Qwen/Qwen3-4B-GGUF:Q8_0 --jinja --color on -ngl 99 -fa on -sm row \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5 \
    -c 40960 -n 32768 --no-context-shift

That is, --color and -fa now needed an on argument.

I kicked it off, and got an error:

Downloading Qwen3-4B-Q8_0.gguf ───────────────────────────────────── 100%


Loading model... \1.54.635.084 E llama_model_load: error loading model: device Vulkan0 does not support split buffers
1.54.635.088 E llama_model_load_from_file_impl: failed to load model
1.54.635.110 E common_fit_params: encountered an error while trying to fit params to free device memory: failed to load model
|1.54.741.517 E llama_model_load: error loading model: device Vulkan0 does not support split buffers
1.54.741.519 E llama_model_load_from_file_impl: failed to load model
1.54.741.521 E cmn  common_init_: failed to load model '/home/giles/.cache/huggingface/hub/models--Qwen--Qwen3-4B-GGUF/snapshots/bc640142c66e1fdd12af0bd68f40445458f3869b/Qwen3-4B-Q8_0.gguf'
1.54.741.523 E srv    load_model: failed to load model, '/home/giles/.cache/huggingface/hub/models--Qwen--Qwen3-4B-GGUF/snapshots/bc640142c66e1fdd12af0bd68f40445458f3869b/Qwen3-4B-Q8_0.gguf'
llama_server exited with code 1
1.54.742.424 E srv  llama_server: exiting due to model loading error
Error: the server exited before becoming ready

That seemed strange. The instructions I'd followed on the Arch Wiki page were the ones to install Llama.cpp for CUDA, but it was referring to device Vulkan0. Vulkan is an open GPU-programming language -- it fits into the stack in essentially the same place as CUDA. As always seems to happen with these things, while the open platform is, well, open, and runs on more platforms than the closed one, it is somewhat slower and buggier.

So, why did I have a Vulkan version? Looking at the talk page for the package was enlightening: user Gotoro wrote:

Hello. The article only provides the llama.cpp-vulkan package for GPU inference, even though on AUR there is a CUDA package available from the same submitter [1]

I must add that the CUDA package is "Flagged out-of-date (2025-12-22)"

...and got a reply from SH3NG1UN saying:

Hi. The reason CUDA was not added is that Vulkan solution is general, performant enough for most hardware setup, and I personally believe that is what edge AI deployment should look like - if your software stack is capable of gaming, it is capable of AI.

They went on to say that they would add CUDA support, but they hit an issue because the AUR version that Gotoro had referred to no longer had an active maintainer. For non-Arch users, the AUR is essentially a repository of non-official, community maintained packages.

So, it looked like I had a Vulkan stack installed. Could I install the CUDA one from the AUR instead?

Over the last few months, there have been security issues with the AUR, where bad actors have picked up ownership of unmaintained packages and put nasty stuff into them -- exfiltration code to steal things like SSH keys and crypto wallets, for example.

As a result, I've been pretty cautious about adding on new AUR dependencies -- and this specific one having been unmaintained for a while scared me.

I decided that I would stick with Vulkan for now; if I wanted to try with CUDA, I'd compile it myself from source later on.

So, that left the error message:

device Vulkan0 does not support split buffers

I decided to break down the command line step by step, as running random stuff you copied from the Internet is rarely a good way to understand what's going on.

llama-cli -hf Qwen/Qwen3-4B-GGUF:Q8_0

We're running llama-cli and specifying a model from Hugging Face -- that's clearly what -hf means.

--jinja

Per the --help, "whether to use jinja template engine for chat".

--color on

"Colorize output to distinguish prompt and user input from generations"

-ngl 99

This is "max. number of layers to store in VRAM". For this model, I wanted everything in VRAM, so felt that "all" would be better there. 99 basically means "all", at least for any normal-sized model, but I wanted to be explicit.

-fa on

This switched on Flash Attention, which seemed reasonable.

-sm row

This was "how to split the model across multiple GPUs". The error was "does not support split buffers", so this sounded relevant! Given that I only have one GPU on my machine, I decided that -sm none would be a better option here.

At this point I suspected I had worked out how to solve the problem, but decided to take a look at the other options just out of interest.

--temp 0.6 --top-k 20

These were clearly the normal sampling parameters that I was used to from creating my own models.

 --top-p 0.95

This was an extension to the previous sampling params: in addition to the top-k stuff (and after it), limit sampling to the tokens that take up the top 0.95 of the probability distribution.

--min-p 0

This was just disabling the min-p feature. min-p discards tokens that are less than X times the probability of the most likely token. Apparently using it is a promising alternative to messing around with top-p and top-k. Rabbit hole alert, let's move on.

--presence-penalty 1.5

Per the docs, this applies a "repeat alpha presence penalty". That was interesting! I googled around a bit and found this page, where it became clear that it is a trick to stop models from getting stuck in loops. Small models are quite prone to that, so it certainly sounded useful. Perhaps something to dig into more later, but for now, keeping it seemed perfectly reasonable.

-c 40960

That was the "size of the prompt context" according to the help -- that is, the amount of the context window that can be taken up by stuff going into the model before it starts generating.

-n 32768

...and this was the other side of that equation: the maximum number of tokens to predict.

--no-context-shift

Another interesting one, "whether to use context shift on infinite text generation". Context shift, it appears, means that if you keep on generating, it will just drop stuff off the start of the sequence when it reaches the max context length, so that it can just keep going. Seems clever, though there are obvious risks of dropping important bits of context (compacting the context would be safer, and that's what most real-world workflows seem to do).

But anyway, we're switching that off, so no harm there.

So, most of those parameters seemed sensible, but I wanted to set -sm to none so that it didn't try to split things over multiple GPUs when I only had one, and also set -ngl to all rather than 99, just for tidiness's sake. That gave this command:

llama-cli -hf Qwen/Qwen3-4B-GGUF:Q8_0 --jinja --color on -ngl all -fa on -sm none \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5 \
    -c 40960 -n 32768 --no-context-shift

...and that worked!

Loading model...

▄▄ ▄▄
██ ██
██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
                                    ██    ██
                                    ▀▀    ▀▀

build      : b10068-571d0d540d
model      : Qwen/Qwen3-4B-GGUF:Q8_0
ftype      : Q8_0
modalities : text

available commands:
  /exit or Ctrl+C     stop or exit
  /regen              regenerate the last response
  /clear              clear the chat history
  /read <file>        add a text file
  /glob <pattern>     add text files using globbing pattern

I gave it a whirl:

> What is your name?

[Start thinking]

Okay, the user is asking for my name. I need to respond with my actual name. Let
me recall, my name is Qwen. I should mention that I'm a large language model
developed by Alibaba Cloud. Also, I should explain that I can help with various
tasks like answering questions, creating content, and more. Keep it friendly and
informative. Make sure the response is clear and concise.
[End thinking]

Hello! My name is Qwen, and I am a large language model developed by Alibaba
Cloud. I can help with a wide range of tasks, such as answering questions,
creating content, and more. How can I assist you today? 😊

[ Prompt: 32.7 t/s | Generation: 122.8 t/s ]

Looking good! Time to try our MoE.

Trying out `Qwen3.6-35B-A3B` -- the first quant

The most popular quants I noticed for this model were Unsloth's, which makes sense -- they're a well-known organisation. They publish their models to Hugging Face, and the model I wanted was there as unsloth/Qwen3.6-35B-A3B-GGUF.

I wasn't sure which of the various 4-bit quants to use -- as I said earlier, there were four of them -- but this Reddit commenter was using UD-Q4_K_XL, so I decided to start with that, just to see what happened.

giles@perry:~ $ llama-cli -hf unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_XL \
    --jinja --color on -ngl all -fa on -sm none \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5 \
    -c 40960 -n 32768 --no-context-shift
Downloading mmproj-BF16.gguf ─────────────────────────────────────── 100%
Downloading Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf ──────────────────────── 100%


Loading model... -ggml_vulkan: Device memory allocation of size 260149504 failed.
ggml_vulkan: vk::Device::allocateMemory: ErrorOutOfDeviceMemory
9.56.395.340 E ggml_gallocr_reserve_n_impl: failed to allocate Vulkan0 buffer of size 260149504


▄▄ ▄▄
██ ██
██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
                                    ██    ██
                                    ▀▀    ▀▀

build      : b10068-571d0d540d
model      : unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_XL
ftype      : Q4_K - Medium
modalities : text, vision, video

available commands:
  /exit or Ctrl+C     stop or exit
  /regen              regenerate the last response
  /clear              clear the chat history
  /read <file>        add a text file
  /glob <pattern>     add text files using globbing pattern
  /image <file>       add an image file
  /video <file>       add a video file



> Hello, who are you?

[Start thinking]

Here's a thinking process:

1.  **Analyze User Input:**
   - User says: "Hello, who are you?"
   - This is a simple greeting and identity question.

2.  **Identify Key Requirements:**
   - Acknowledge the greeting
   - State my identity clearly and concisely
   - Mention my capabilities/purpose briefly
   - Keep it friendly and helpful

3.  **Formulate Response (Internal Refinement):**
   - "Hello! I'm Qwen, a large language model developed by Alibaba Group's
     Tongyi Lab. I'm here to help with answering questions, writing, coding,
     analysis, and much more. How can I assist you today?"
   - Check against guidelines: Matches identity, concise, friendly, offers help.

4.  **Final Output Generation:** (matches the refined version)✅
[End thinking]

Hello! I'm Qwen, a large language model developed by Alibaba Group's Tongyi Lab.
I'm here to help with answering questions, writing, coding, analysis,
brainstorming, and much more. How can I assist you today?

[ Prompt: 9.8 t/s | Generation: 111.4 t/s ]

That warning at the start was a bit concerning:

Loading model... -ggml_vulkan: Device memory allocation of size 260149504 failed.
ggml_vulkan: vk::Device::allocateMemory: ErrorOutOfDeviceMemory
9.56.395.340 E ggml_gallocr_reserve_n_impl: failed to allocate Vulkan0 buffer of size 260149504

...and I noticed that in nvtop, my VRAM was maxed out.

I decided to remove the -c 40960 -n 32768 from the command line; these were setting a specific context length, both for the prompt and for generation. Llama.cpp can work out what context length it can fit into VRAM on its own if you don't specify them, and this auto-fit felt like a good idea given that I was trying to jam a 22.4 GiB model into a 24 GiB GPU, and had X and various other apps running on the card at the same time.

So:

giles@perry:~ $ llama-cli -hf unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_XL \
    --jinja --color on -ngl all -fa on -sm none \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5 \
    --no-context-shift


Loading model...

▄▄ ▄▄
██ ██
██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
                                    ██    ██
                                    ▀▀    ▀▀

build      : b10068-571d0d540d
model      : unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_XL
ftype      : Q4_K - Medium
modalities : text, vision, video

available commands:
  /exit or Ctrl+C     stop or exit
  /regen              regenerate the last response
  /clear              clear the chat history
  /read <file>        add a text file
  /glob <pattern>     add text files using globbing pattern
  /image <file>       add an image file
  /video <file>       add a video file



> Hello, who are you?

[Start thinking]

Here's a thinking process:

1.  **Analyze User Input:**
   - User says: "Hello, who are you?"
   - This is a simple greeting and identity question.

2.  **Identify Key Requirements:**
   - Acknowledge the greeting
   - State my identity clearly
   - Keep it concise and friendly
   - Align with my actual identity (Qwen / Tongyi Qianwen, developed by Alibaba Group's Tongyi Lab)

3.  **Formulate Response:**
   - Start with a greeting back
   - State name/identity: I'm Qwen (or Tongyi Qianwen), a large language model developed by Alibaba Group's Tongyi Lab
   - Offer assistance
   - Keep it natural and concise

   Draft: "Hello! I'm Qwen, a large language model developed by Alibaba Group's Tongyi Lab. How can I assist you today?"

4.  **Self-Correction/Verification:**
   - Matches identity guidelines? Yes.
   - Concise and direct? Yes.
   - No extra fluff? Yes.
   - Ready.✅
[End thinking]

Hello! I'm Qwen, a large language model developed by Alibaba Group's Tongyi Lab. How can I assist you today?

[ Prompt: 50.2 t/s | Generation: 122.0 t/s ]

No error, it wasn't maxing out VRAM, and generation and prompt processing were noticeably faster. Perhaps the first time around it had offloaded some of the model to the CPU so that it had enough space on the GPU for the context size I was asking for?

Still, I wasn't sure how far I could trust those numbers -- the number of tokens being emitted was quite small. ⁴

I wanted a prompt that would make the model generate lots of tokens. Now, it was a thinking model, so asking it a difficult question seemed like a good way to do that.

I decided to use Ethan Mollick's "Lem test":

Compose a poem -- a poem about a haircut! But lofty, tragic, timeless, full of love, treachery, retribution, quiet heroism in the face of certain doom! Six lines, cleverly rhymed, and every word beginning with the letter S!

More about that in the appendix below. ⁵

Unfortunately, when I gave it that prompt, it thought for a while, but crapped out during the thinking process, while it was trying to generate rhymes.

Now, I didn't know what the context length was, either for the prompt or for generation. But I could be fairly sure that they were less than 40,960 and 32,768 respectively, given that the model did not fit into VRAM when I specified those limits, but did when I asked Llama.cpp to auto-fit.

I couldn't find a quick and easy way of asking llama-cli what lengths it had fit things to, but when I tried running it with the -v option for verbose logging, I got something. Actually, I got a lot of something -- thousands of lines of output, with incredibly detailed messages about everything Llama.cpp was doing.

But there were three blocks starting "constructing llama_context".

...
0.00.809.283 I llama_context: constructing llama_context
0.00.809.285 I llama_context: n_seq_max     = 1
0.00.809.285 I llama_context: n_ctx         = 262144
0.00.809.285 I llama_context: n_ctx_seq     = 262144
0.00.809.285 I llama_context: n_batch       = 2048
0.00.809.285 I llama_context: n_ubatch      = 512
0.00.809.285 I llama_context: causal_attn   = 1
0.00.809.286 I llama_context: flash_attn    = enabled
0.00.809.286 I llama_context: kv_unified    = false
0.00.809.287 I llama_context: freq_base     = 10000000.0
0.00.809.287 I llama_context: freq_scale    = 1
0.00.809.287 I llama_context: n_rs_seq      = 0
0.00.809.287 I llama_context: n_outputs_max = 1
...
0.01.029.515 I llama_context: constructing llama_context
0.01.029.517 I llama_context: n_seq_max     = 1
0.01.029.517 I llama_context: n_ctx         = 4096
0.01.029.517 I llama_context: n_ctx_seq     = 4096
0.01.029.517 I llama_context: n_batch       = 2048
0.01.029.517 I llama_context: n_ubatch      = 512
0.01.029.517 I llama_context: causal_attn   = 1
0.01.029.517 I llama_context: flash_attn    = enabled
0.01.029.518 I llama_context: kv_unified    = false
0.01.029.518 I llama_context: freq_base     = 10000000.0
0.01.029.519 I llama_context: freq_scale    = 1
0.01.029.519 I llama_context: n_rs_seq      = 0
0.01.029.519 I llama_context: n_outputs_max = 1
0.01.029.519 I llama_context: n_ctx_seq (4096) < n_ctx_train (262144) -- the full capacity of the model will not be utilized
...
0.06.393.122 I llama_context: constructing llama_context
0.06.393.123 I llama_context: n_seq_max     = 1
0.06.393.123 I llama_context: n_ctx         = 4096
0.06.393.123 I llama_context: n_ctx_seq     = 4096
0.06.393.123 I llama_context: n_batch       = 2048
0.06.393.123 I llama_context: n_ubatch      = 512
0.06.393.123 I llama_context: causal_attn   = 1
0.06.393.124 I llama_context: flash_attn    = enabled
0.06.393.124 I llama_context: kv_unified    = false
0.06.393.126 I llama_context: freq_base     = 10000000.0
0.06.393.126 I llama_context: freq_scale    = 1
0.06.393.126 I llama_context: n_rs_seq      = 0
0.06.393.126 I llama_context: n_outputs_max = 1
0.06.393.127 I llama_context: n_ctx_seq (4096) < n_ctx_train (262144) -- the full capacity of the model will not be utilized
...

That made the issue reasonably clear. I assumed that the last message was the definitive one, and it was saying that my context length -- both prompt and generation combined -- was 4,096 tokens. I had --no-context-shift set, so that was a hard limit. And 4k or so tokens fit with the amount of stuff that I would have expected to be in the model's context window at the time it crapped out.

I decided to see what would happen if I tried a smaller quant.

Trying out `Qwen3.6-35B-A3B` -- the second quant

Hugging Face said (and nvtop agreed) that this one, UD-Q4_K_XL, was about 22.4 GiB in size. But UD-IQ4_NL_XL was shown on HF as using 19.5 GiB, so I decided to give that one a go next.

I ran it with the -v option, and again looked at the last "constructing llama_context" block:

42.33.956.367 I llama_context: constructing llama_context
42.33.956.369 I llama_context: n_seq_max     = 1
42.33.956.369 I llama_context: n_ctx         = 93952
42.33.956.369 I llama_context: n_ctx_seq     = 93952
42.33.956.369 I llama_context: n_batch       = 2048
42.33.956.369 I llama_context: n_ubatch      = 512
42.33.956.369 I llama_context: causal_attn   = 1
42.33.956.370 I llama_context: flash_attn    = enabled
42.33.956.370 I llama_context: kv_unified    = false
42.33.956.373 I llama_context: freq_base     = 10000000.0
42.33.956.373 I llama_context: freq_scale    = 1
42.33.956.373 I llama_context: n_rs_seq      = 0
42.33.956.373 I llama_context: n_outputs_max = 1
42.33.956.374 I llama_context: n_ctx_seq (93952) < n_ctx_train (262144) -- the full capacity of the model will not be utilized

93,952 tokens -- still not the full amount, but not at all bad! I decided to give it a whirl, restarted without -v to avoid getting swamped with debug info, and:

> Compose a poem -- a poem about a haircut! But lofty, tragic, timeless, full of love, treachery, retribution, quiet heroism in the face of certain doom! Six lines, cleverly rhymed, and every word beginning with the letter S!

[Start thinking]

Thinking Process:

1.  **Deconstruct the Request:**
    *   Topic: Haircut.
    *   Tone: Lofty, tragic, timeless, full of love, treachery, retribution, quiet heroism in the face of certain doom.
    *   Length: Six lines.
    *   Rhyme: Cleverly rhymed.
    *   Constraint: Every word must begin with the letter 'S'.

2.  **Brainstorming Vocabulary (S-words):**
    *   *Nouns:* Shears, scissors, scalp, skin, strand, silk, sorrow, shame, secret, sister, son, star, storm, shadow, silence, sacrifice, splendor, silver, sword, slaughter, severance, sorrow, sigh.
    *   *Verbs:* Sever, slice, shatter, shed, sink, slip, slide, sing, scream, suffer, survive, surrender, slay, snatch, steal, spin, sway, smile, stare, stand, seek, save, seal, send, set, show, shut, sing, sit, sleep, slip, slow, snap, sneak, soar, speak, spend, spill, split, spoil, spread, spring, squeeze, stand, start, stay, steal, stem, stir, stop, store, stretch, strike, string, strip, strive, stroke, strong, study, suffer, suggest, suit, sum, sun, surprise, surround, survey, survive, suspect, sustain, swallow, swear, sweep, sweet, swell, swim, swing, switch, symbol, sympathize, system.
    *   *Adjectives/Adverbs:* Silent, silver, sad, severe, soft, sharp, slick, slim, slow, small, strange, strong, sweet, swift, sure, sudden, supreme, sacred, secret, solemn, sorrowful, splendid, stark, still, stern, straight, true.

3.  **Drafting - Attempt 1 (Focusing on Narrative and Tone):**
    *   *I:* Need love, treachery, retribution, heroism.
    *   *I:* Lover betrays? Or self-sacrifice? "Quiet heroism in the face of certain doom" suggests a sacrifice. Maybe the hair is the doom? Or the haircut is the retribution?
    *   *Concept:* A queen/heroine sacrifices her beauty/hair to save someone or fulfill a vow, betrayed by the scissors or the situation.
    *   *Lines:*
        *   Sweet shears slice silver strands... (Too simple).
        *   She saw sharp steel... (Better).
        *   Secret sorrow strikes...
        *   Let's try rhymes. AABBCC or ABABCC or ABCABC? "Cleverly rhymed" might suggest internal rhyme or a complex scheme, but with 6 lines and all S-words, the rhyme scheme needs to be tight. Maybe AAAA AA? Or AAB BCC?

[LLM pondering snipped]

        *   Result matches constraints.
        *   Six lines. All S words. Rhyme A B C B C C (Share/Snare, Soul/Sole, Son/Sun/Son).
        *   Share/Snare is exact.
        *   Soul/Sole is exact (in many dialects, or slant in others, but usually accepted as rhyme with Sold/Soul/Sole cluster).
        *   Actually, Soul/Sole can be distinct.
        *   Soul is /soʊl/. Sole is /soʊl/. They are homophones in many accents.
        *   So Soul/Sole is exact rhyme.
        *   Son/Sun is slant.
        *   This is a strong poem.
[End thinking]

Sneering shears shred sacred share,
Silent steel sells silver soul;
Sister's snare spares son;
Sweet sign seals sole;
Still star saves sun;
Secret sign seals son.

[ Prompt: 183.3 t/s | Generation: 119.6 t/s ]

It's hardly Shakespeare, and the rhyming is both odd (ABCBCC) and unadventurous ("son", "sun", "son") but for a 35B parameter model with 3B active (of which 1B are used up by the embeddings and the output head), it's pretty bloody impressive! And every word did indeed start with an "S".

So, now I had the model running. I had a context length of 93,952 tokens, was generating about 120 tokens/second, and processing the prompt at about 180 tokens/second.

The question was, could I get the full context length in there? I knew that the answer was yes from previous reading, but it was time to try it out.

The `-ncmoe` flag

Llama.cpp has a command-line flag -ncmoe, or in its long form, --n-cpu-moe. As its name suggests, it's designed for mixture of experts models, and it works by offloading the FFN for a specified number of Transformers layers to the CPU.

Intuitively it's surprising that that doesn't cause performance to completely crash. After all, aren't we running our models on the GPU because it's so much better at matrix multiplications than the CPU?

And it's true, the GPU is best. But if it's a question of where to use its superiority in matmuls, then keeping the attention layers on the GPU and offloading just the FFNs to the CPU is the least harmful way to do it. Obviously you do lose performance, but not as much as you would if you offloaded attention.

I started a binary chop. From the logs, I could see that the model had 40 layers:

0.00.804.467 I print_info: n_layer               = 40
0.00.804.467 I print_info: n_layer_all           = 40

so I ran this to offload the first 20 to the CPU:

llama-cli --log-file /tmp/log --log-colors off -hf unsloth/Qwen3.6-35B-A3B-GGUF:UD-IQ4_NL_XL \
    -ncmoe 20 --jinja --color on -ngl all -fa on -sm none \
    --temp 0.6 --top-k 20 --top-p 0.95 --min-p 0 --presence-penalty 1.5
    --no-context-shift

That gave me a context length of 262,144 -- the native context length! I had 16944MiB used on the GPU, 10075MiB on the CPU, and I got 51.4 tokens per second for generation, and 64.4 tokens per second for the prompt.

Then I tried various other values, but the process got dull quickly. For each one, I was doing this:

Start llama-cli with the -ncmoe value I wanted and with -v. Use the logs to find the context length, and note down the VRAM/RAM usage.
Restart it without -v, paste in my prompt, and wait until it either came back or got stuck in a loop trying to find rhymes (which it did maybe one in every four times). Note down tokens per second for the prompt and generation.

At maybe four minutes for each and an error-prone process, this was ripe for automation.

Additionally, something about those numbers for the prompt-processing tokens per second was niggling me. They seemed very slow. You'd expect an LLM in a generation harness like this to process the prompt much faster than it could generate, but I was getting numbers that roughly matched the generation speed. I figured that my prompt was so short that the numbers I was getting were being dominated by overhead. I needed a longer one.

The auto-sweep

Llama.cpp's llama-server command exposes an API that tells you what sequence length it's running with, and takes pretty much the same parameters as llama-cli, so I gave Claude Fable 5 the command I'd been running and asked it to write a script to sweep over CPU offload layer counts from zero to 40, and to store the numbers I was interested in:

The number of layers offloaded
The context length we got
The VRAM and system RAM (RSS) usage
The size of the prompt we were providing
The tokens per second at which the prompt was processed
How much was generated
Generation tokens/second.

I also added onto that some code to expand the prompt so that I could get better numbers for the tokens per second for that part of the processing.

Now, llama-cli uses a chat template by default; the endpoint that Claude had chosen for llama-server does not -- it's a raw completion engine. This actually worked to my advantage. From the llama-cli logs I found the template it had been using, which was something like this:

'<|im_start|>system
You are a helpful assistant<|im_end|>
<|im_start|>user
Hello<|im_end|>
<|im_start|>assistant
Hi there<|im_end|>
<|im_start|>user
How are you?<|im_end|>
<|im_start|>assistant
<think>
'

Because I could easily fake the conversation history, I baked in a template that started with a fake previous interaction where the user had provided the first two chapters of Jane Austen's Pride and Prejudice and asked the model for its opinion, it said something like "I like that", and then the user asked for the Lem test poem. That bulked things out quite a bit -- as it turned out, to a prompt of 2,457 tokens.

I also made it give up after 6144 tokens so that it wouldn't take forever to run, and kicked it off. Here's the code, and you can see the results here.

The prompt processing numbers were much better! With nothing offloaded, I got a context window of 51,200, a prompt processing rate of 2,787.1 tokens/second, and a generation rate of 122.4 tokens/second.

(It was interesting that the context window was smaller than it was when I ran llama-cli without -ncmoe, but I decided that I didn't want to get into the weeds by digging into that ⁶. Perhaps -ncmoe 0 isn't quite the same as running without -ncmoe at all?)

Anyway, the numbers for other levels of layer-offloading did what you would expect: as more layers' FFNs went from the GPU to the CPU, the context window expanded until it hit the model's native context length of 262,144 when 12 layers were offloaded. And the speed -- both in prompt processing and generation -- went down.

I ran it three more times, just to see if there was a lot of noise there, and then asked Claude Sonnet to slop up a charting script for me (in the same repo as the above if you're interested) -- and I had some results!

The results

Just for people who skipped the details: we're running the UD-IQ4_NL_XL quant of unsloth/Qwen3.6-35B-A3B-GGUF on llama-server, and seeing how its performance changes as we vary the number of layers offloaded from the GPU to the CPU.

Firstly, how does the context length change?

You can see that the model's full built-in context length was reached when we had 12 layers offloaded. But how did that impact performance?

You can see that the prompt throughput drops off fairly smoothly as the number of layers offloaded increased. What is interesting is those spikes in the generation tokens/second. The reason I'd done four runs of this script was because I'd initially thought they were noise, but you can see they're pretty consistent.

Claude Fable 5 thinks they may be related to the points at which there are full attention layers -- the attention system in these Qwen models is complex, with some layers using a faster system and others using the full attention that my GPT-2 models always use. That's something to investigate another time, I suspect, if ever.

It's noticeable that with 12 layers offloaded (the exact number that gave us the full context), the throughput is somewhat better than at 11, which makes 12 look pretty much like the sweet spot for running this model with this setup! From the raw numbers, we can see that there we had about 66 tokens/second generation, and 607 tokens/second on the prompt.

Anyway, finally, let's take a look at memory usage -- both VRAM and regular system RAM. It looks exactly as you'd expect:

We're using as much VRAM as we can get up to the layer 12 offload point, because with fewer than that number of layers offloaded, we're grabbing as much space as we can on top of the amount consumed by the parameters, in order to get the longest possible context window. But after that, it drops off smoothly. And system RAM rises smoothly as layers are offloaded into it, as expected.

So, that was it!

Or was it?

CUDA

As I mentioned back at the start of these lab notes, the version of Llama.cpp that's part of the official OS repositories for Arch uses Vulkan, in order to maximise the number of platforms it runs on. And that was the version I'd used.

The Vulkan version was clearly solid and stable, given that I'd been able to run a bunch of benchmarks on it. But you'd expect CUDA to get performance enhancements and the like first, and in general to be more polished.

So how would it perform comparatively?

I didn't want to install the AUR version of the CUDA-based Llama.cpp, as it had been abandoned previously (though someone appears to have picked it up again now), and bad actors have been taking over abandoned AUR repos and putting bad stuff in them.

So out of an abundance of caution, I decided to compile it from source from the official Llama.cpp repo. Their instructions made it almost laughably simple; I cloned it, then ran

cmake -B build-cuda -DGGML_CUDA=ON -DGGML_NATIVE=ON
cmake --build build-cuda --config Release -j$(nproc)

Less than ten minutes later, I had a working set of binaries. I tweaked my benchmarking script to use them instead of the system-installed ones, and got these results.

Charting those, we get this for the context length:

So with CUDA, we get the full built-in context length with just 10 layers offloaded, compared to the 12 we needed with Vulkan.

And even with nothing offloaded, we got a bigger context window -- 89,600 rather than Vulkan's 51,200.

How about performance?

With CUDA, we don't get those odd bumps in the generation throughput -- or at least, if they're there, they're much smaller. Perhaps Claude's ponderings that I mentioned earlier were off the mark? Or perhaps they only impact Vulkan? ⁷

But the big news on this chart is -- as I had suspected would be the case -- CUDA was noticeably faster! With all layers on the GPU, it was getting 139.6 tokens/second generation, 3360.4 tokens/second on the prompt, as compared to around 122 / 2787 for Vulkan, averaging across my four runs on that platform. And with a bigger context window too.

With ten layers offloaded, CUDA had the full context length of 262,144, and was getting 89.1 / 1153.5, while Vulkan managed a context length of 233,984, and throughput of 66 / 702.

Finally, with 12 layers offloaded, both models had the full context length, CUDA had throughput of 84.9 / 1008.1, and Vulkan had 66 / 607. And that meant that we were able to reclaim some VRAM too:

Advantage: CUDA, I think.

Conclusion

That was probably just a tad more time and effort than I think my friend expected me to put into this when he asked an offhand question on WhatsApp. But I feel that I've learned a lot of useful stuff in this journey, and while there's a lot of detail in this post, it didn't actually take up much time to run the experiments. I have another post on the way that involves multiple four-day training runs, so I needed to occupy my time somehow...

Anyway, I hope it's useful to people out there! Let me know in the comments if it helped you.

And now, just to finish off, a bonus section on the prompt I was using to get the model to generate lots of stuff.

Appendix: the Lem test

The prompt I was using for these tests was this:

Compose a poem -- a poem about a haircut! But lofty, tragic, timeless, full of love, treachery, retribution, quiet heroism in the face of certain doom! Six lines, cleverly rhymed, and every word beginning with the letter S!

What I like about it as a prompt is not just that it's a particularly silly thing to ask of an LLM, but it forces models with reasoning enabled to think hard about how to keep within all of those constraints. It's basically a short prompt that makes them generate a lot of tokens -- perfect for benchmarking.

I stole it from Ethan Mollick. He's well worth following if you're not already doing so; here's his X/Twitter profile, and his Substack. He writes a ton of interesting stuff, but I want to focus here on this particular eval.

In the science fiction short story "The first sally (A), or Trurl's electronic bard" (included in the collection The Cyberiad), Stanisław Lem wrote of two rival engineers. One of them, Trurl, has created an electronic poet, and the other, Klapaucius, gives it what he thinks is an impossible task. From the English translation of this part on Goodreads:

"Have it compose a poem -- a poem about a haircut! But lofty, tragic, timeless, full of love, treachery, retribution, quiet heroism in the face of certain doom! Six lines, cleverly rhymed, and every word beginning with the letter S!!"

The machine ponders briefly, and responds:

Seduced, shaggy Samson snored.
She scissored short. Sorely shorn,
Soon shackled slave, Samson sighed,
Silently scheming,
Sightlessly seeking
Some savage, spectacular suicide.

The story is a fun read, definitely recommended! You can see a bunch of parallels between current LLMs in the way Trurl's machine works -- and in its failure modes.

Mollick decided to use it as a prompt ages ago, and in March 2024, he managed to get Claude 3 to almost complete the task. The result wasn't a very good poem, but neither ChatGPT 4 nor the version of Gemini at the time were close. Here's Claude 3's attempt:

Stealthily snipping, shearing silky strands,
Severing soulmate's sacred, sworn soulbands.
Scoundrel's scandalous sins spurred sly scheme,
Spurring scorned sweetheart's spiteful supreme dream.
Shining locks sacrificed, stoically shorn,
Stalwart spirit stays strong, staunchly reborn.

You can see that it had two non-S words, "dream" and "reborn" ⁸. But otherwise not bad -- it makes sense, the rhymes work, and it ticks the emotional boxes.

Recent models have become very good; here's Claude Fable 5 from today:

Ah, Trurl's Electronic Bard's finest hour — the impossible commission from The Cyberiad! Here's my own attempt at the feat:

Soft she sang; strong Samson slumbered, serene;
Sly, she snipped seven shining strands, scarcely seen.
Soldiers stormed, seized, scorched Samson's sight;
Shackled, still, sightless — seething sacred spite.
Straining, steadfast, Samson seized stout stanchions, swaying slow —
Stones smashed; seductress, soldiers, Samson: slain. So sagas show.

Love in the lullaby, treachery in the shears, doom in the darkness, heroism at the pillars, and retribution in the rubble — with every single word dutifully saluting the letter S. The Electronic Bard would, I hope, give a grudging beep of approval.

Now, it recognised the source, so perhaps that helped -- but while I'm no connoisseur of poetry, I actually think it's slightly better than Lem's version (though as Lem wrote in Polish, that's the work of his translator -- and translating poetry is obviously a tough task).

And while both versions are based on Samson, given that he was a biblical hero who lost his strength after his hair was cut, he's a pretty obvious subject for a tragic poem about a haircut, especially if every word has to start with the letter "S" -- so I don't think that points to copying.

ChatGPT 5.6 Sol also does a great job:

Silken Samson sought sweetheart’s soft sanctuary;
Sweetheart, serpent-souled, summoned shears—sanguinary.
Silver strands slid; sacred strength sank, slain;
Silent sorrow sealed sweetheart’s scarlet stain.
Shorn, storm-bound, Samson shielded slaves, steadfast;
Stone sanctums shattered; smiling, Samson sank—surpassed.

Another Samson-based one; it doesn't say it recognises the source, but I imagine it did.

Now let's take a look at some samples from this tiny model we've been benchmarking -- 35B parameters, 3B active, of which only 2B are doing the actual thinking (rather than messing around with embeddings). Here's the first one I got:

Sneering shears shred sacred share,
Silent steel sells silver soul;
Sister's snare spares son;
Sweet sign seals sole;
Still star saves sun;
Secret sign seals son.

I mean, it's not amazing, but it satisfies the S constraint. Less impressed by how well it handles the "cleverly rhymed" part, though -- as I said earlier, its scheme is ABCBCC, which is odd rather than clever -- and it doesn't make an awful amount of sense.

I noted down two more while I was playing with this:

Silent shears sever sweet silver strands;
She stands, sacrificing self; submission seals severing's demands.
Shear's snare strikes; swift sentence seals sorrowful sphere.
Shorn sovereign swears silence shields secret spear;
She saves sacred soul; stillness safeguards sterling store.
Spirit soars; sorrow's stroke stops, solace restores.

That one's a lot better, I think; there's a hint of a story, and at least it rhymes. But note the "demands" and the "restores" at the ends of lines 2 and 6 -- the "S" constraint has been broken.

The next:

Sister's scheme scalp shorn,
Silver shears slice soul sworn,
Sweetheart's steel skull sear,
Silent slitters shear.
Spirit spies sight,
Sorrow swallows slight.

Hmm. A bit bleak.

Again, poetry isn't my thing. But I think that Claude 3's offering from two years ago is better than any of them -- apart from the fact that it couldn't keep the "S" constraint.

Still, it's impressive to see how well such a tiny model -- again, 2B actual thinking parameters per token -- can do on what is a pretty challenging task!

With some models, there are also extra FFNs that are always loaded that do stuff before the routed expert ones. Again, I'll dig into the details at a later point. ↩
That said, a while back Dan Woods managed to get Qwen3.5-397B-A17B -- 397B parameters with 17B active -- to run on a 48 GiB Macbook Pro with some very clever use of the SSD. It was slow, though, at around 4 tokens/second. ↩
And only supports certain sizes for certain kinds of operations. ↩
As was the size of the prompt, but I got to that later. ↩
Yes, my blog posts have started accumulating appendices as well as footnotes. I may have a problem. ↩
Some readers might feel that that ship sailed long ago and I was already very much into the weeds, and possibly in the middle of the Amazon jungle. I disagree, but largely due to the mixed metaphor. ↩
"Put down that rabbit hole and step slowly away, Giles." ↩
'Hey, Claude 3, how many "S"s are there in "strawberry"?' ↩

Why do OpenAI's GPT-2 weights beat mine?

Wed, 29 Jul 2026 15:30:00 +0000

When I finished my project training an LLM from scratch, I was left with a minor mystery. Why were my models worse at instruction-following than the original OpenAI GPT-2 small weights?

I had an evaluation that I was running, based on the instruction fine-tuning code in chapter 7 of "Build a Large Language Model (from Scratch)". The process was to train a model on samples from the Alpaca instruction-following dataset until validation loss started rising, to use that instruction fine-tuned model to generate completions to a held-back test set, and then to use an LLM to compare the results from various different models. The details are here; let's call it the IFT eval.

OpenAI's original weights for GPT-2 small consistently beat my own models, even when mine got better results than theirs on a more technical evaluation, where I just measured the cross entropy loss for each model on a held-back set of test sequences. This surprised me; I would have expected a reasonably close correlation between the two evals -- that better test loss would imply better instruction-following.

I have a couple of thoughts about why this might be, and given that I recently set up poppy, my dedicated LLM training box I decided to inaugurate her with an experiment to test one of them; further experiments will come in time -- though I don't think this will be a focus for the blog. More of a running theme, with occasional posts until either I solve the mystery, or give up in despair...

In this post I'll give a bit more detail about the nature of the problem, and list some of the things I've been thinking might be the cause. In later posts, I'll dig into some of them.

More on the mystery

Let's take a look at the results from my most recent runs of the IFT eval. I've highlighted the OpenAI models in bold, and the table is sorted by the test loss -- that more technical evaluation that I mentioned earlier, where lower is better.

How well the model did with the IFT eval is in the last three columns. The "IFT epochs" column shows how many epochs of training the model needed before its validation loss started rising, the score is the average mark out of 100 that (in this case) GPT 5.5 gives the model's answers to the IFT test set, and the rank is the position the model holds in terms of that average score.

	Test loss	IFT epochs	IFT score	IFT rank
OpenAI weights: medium	3.231442	2	41.62	1
JAX, with MHA bias, no dropout	3.418784	4	19.25	4
JAX, no MHA bias, no dropout	3.420089	3	14.66	11
JAX, no MHA bias, with dropout	3.476802	4	12.94	15
OpenAI weights: small	3.499677	2	26.73	2
`1xrtx3090-stacked-interventions`	3.538161	4	17.79	6
`8xa100m40-stacked-interventions-1`	3.577761	4	10.29	16
Cloud FineWeb, 8x A100 40 GiB	3.673623	7	20.71	3
`1xrtx3090-baseline`	3.683835	6	15.11	9
`8xa100m40-baseline`	3.691526	4	14.74	10
Cloud FineWeb, 8x H100 80 GiB	3.724507	4	13.25	14
Cloud FineWeb, 8x A100 80 GiB	3.729900	4	14.50	12
Cloud FineWeb, 8x B200 160 GiB	3.771478	4	16.03	8
Local FineWeb train	3.943522	7	13.73	13
Local FineWeb-Edu extended train	4.134991	7	16.70	7
Local FineWeb-Edu train	4.166892	7	18.68	5

For the IFT eval, the OpenAI weights are the best. The "medium" model comes first, and the "small" one comes second -- and there's a noticeable gap between "small" and the score for my best model for that eval, "Cloud FineWeb, 8x A100 40 GiB": 26.73 vs 20.71. That difference was consistent across multiple runs of the eval -- the relative rankings of my own models varied more.

It's not surprising that the "medium" model does better -- it's more than double the size of the others -- but the small model consistently beating mine was odd given that I had models with lower test loss.

Another thing that stands out is the number of epochs that the models were trained for. Remember, I trained them until the validation loss started rising -- that is, until the model started overfitting. You can see that the OpenAI weights hit that after two epochs of training, whereas my models ranged between three and seven.

The loss landscape

My starting hypothesis was that the OpenAI weights were starting off from a better position in the IFT loss landscape than mine. The loss landscape for the task the various models were originally pre-trained for -- "predict the next token for this sequence of text that was pulled from the web" -- is different to the loss landscape for the IFT task, which is something more like "predict the next token in a useful response to this request". Let's dig into that a bit.

The core idea behind the normal LLM training paradigm is that we can get a useful starting point for a model by pre-training -- training on a whole load of cheaply-available stuff from the web -- and only then have to use our much more expensive task-specific datasets for fine-tuning. I've been training my models on approximately 3.2 billion tokens of the fineweb dataset, which was generated from scrapes of the web, and then deduplicated and tidied up a bit.

If I had instead trained them on 3.2 billion tokens of instruction-following samples -- essentially, pre-packaged Q&A sessions between a human and an assistant -- then they would almost certainly be much better than the models I have at instruction-following, for the same cost in compute used to do the training run.

But the cost of generating those 3.2 billion tokens would be incredibly high. Imagine hiring people to do it: the GPT-2 tokeniser averages about 0.75 words per token, and while I don't know how much you'd need to pay people to write 2.4 billion words, I suspect it would be rather a lot. Even generating that much synthetic data from a larger LLM (as the Alpaca dataset was) would not be cheap -- I make it US$144,000 using GPT 5.6 Sol as of mid-2026.

Even worse, the LLM we trained on that kind of data would be "brittle" -- if you wanted it to solve an even slightly different task (for example, Alpaca is single-shot question and answer, so imagine if you wanted it to handle multi-turn conversations like a chatbot), it would be hard to train it to do that.

So: the idea is that if we pre-train a base LLM on less-structured -- but, importantly, still "real" -- data like scraped web pages, we'll get a general-purpose base LLM relatively cheaply. It will have discovered basic stuff like the structure of language, and hopefully some general knowledge (eg. maybe that the capital of France is Paris). Once we have that, we can fine-tune it for specific uses.

Now, both the initial pre-training and -- in this IFT test -- the fine-tuning phase of building our model are done in essentially the same way: trying to minimise the cross entropy loss of the model's predictions against a dataset. ¹ What has changed between them is the targets we're trying to train the model to predict. In the pre-training phase, we're trying to get it to predict web scrape data, in the fine-tuning, we're trying to predict responses to queries.

Putting that idea into slightly more mathematical language (albeit loosely enough that any real mathematicians reading this will probably scream with horror), if we're saying that a pre-trained model is useful as a base for instruction fine-tuning, what we have is an assumption that the loss landscape for the original pre-training phase is reasonably similar to the loss landscape for the fine-tuning. We are hoping that a place that is nice and low on the pre-training landscape is reasonably close (in parameter space) to a place that is low on the fine-tuning landscape.

That's all very abstract; let's try to visualise it. If we imagine the loss landscape for a model with two parameters, that's reasonably easy. The two parameters make up two dimensions -- let's say left and right for one, forward and back for the other -- and the loss at any given point is the vertical dimension. It's an uneven surface, a bit like a rolling landscape, with hills and mountains at points where the parameters have very high loss, and valleys and clefts where the loss is lower.

In reality, of course, we have somewhat more than two parameters. With the 163,009,536 parameters in my GPT-2-small-style models, the loss landscape is not a surface but a 163,009,536-dimensional hypersurface ² in 163,009,537-dimensional space. Good luck visualising that.

While taking intuitions from levels of dimensionality that we can actually imagine over to insanely high-dimensional spaces like that is often risky -- weird stuff starts happening as the number of dimensions goes up -- for what we're specifically looking at here, it's safe. The landscape image works for intuition.

So, we have one landscape like that for our pre-training phase. When we start fine-tuning, the loss landscape changes to a different one. What we're hoping is that the landscape is reasonably similar; that the low point that we wound up at in the pre-training landscape is close to a low point in the new fine-tuning landscape when that is swapped in. ³

The good news is that we know that this process of doing a pre-train then a fine-tune works. Most modern AI is trained using it as a foundation (though there's lots of extra stuff on top). And indeed, you'd intuitively expect it to work -- it would be kind of weird if there was no correlation at all between the "understand language" loss landscape and the "answer questions" one.

But this does mean that we can -- at least in a somewhat hand-wavey way -- characterise the manner in which the OpenAI weights are "better" than mine in terms of the loss landscape.

Both the OpenAI weights and mine have landed in places that were good from the pre-training viewpoint; that's what the "Test loss" results in the table above mean. GPT-2 medium is doing better than any of my models -- given that it's about twice the size, it should -- and GPT-2 small is doing better than most of them.

But the places in the pre-train loss landscape where the OpenAI models landed were clearly better-matched to good places in the fine-tuning loss landscape than my own models' places were. The fact that they took fewer epochs to start overfitting intuitively points in that direction -- after all, if you're closer to somewhere, then it takes less time to get there -- but the fact that they scored better after fine-tuning makes it pretty clear.

But something else that occurs to me as I write this is that the results on the test set also point to a superiority in the OpenAI weights -- one that hadn't occurred to me in the past.

OpenAI weights: better than we thought!

Let's think about what that test set is. In order to get data that was pre-processed and ready to work with, I downloaded all of the 10 billion token version of FineWeb, and split it into 99% training and 1% validation and test. I tokenised each sample in each of those splits, and then concatenated them together into a single sequence for each, separating the samples with <|endoftext|> tokens. My training script took the train split, broke it up into 1024-token sequences, assembled them into batches, and ran through 3.2B token's worth.

The loss test takes 19,660,800 tokens starting at position 50,000,000 in the validation/test split, and runs them through in batches of six, working out the average loss.

Now, the GPT-2 weights were trained on something that I believe was constructed similarly -- get a load of text consisting of a bunch of documents from the web, tokenise them all and tack them together separated with <|endoftext|>s, and then use that.

But, importantly, it was a different dataset. Annoyingly, we don't have access to OpenAI's "WebText" dataset, but I think it's reasonably safe to say that even if it is similar to what I came up with, it will be less similar to the sequences used for the loss test than my own training set is.

Given that the loss test for the OpenAI small weights came out with a better result than any of the models I trained in PyTorch, and was only narrowly beaten by the ones I trained with JAX (even then, I think, because the JAX models got lucky with their random weight initialisation before they started training), then I think we can say that the OpenAI weights are already considerably better than my own ones. They were already at a good place in their own loss landscape, but it happened to also be in a good place in my test set's.

An alternative analogy: if you imagine a group of 15 trail runners having a race on a route in a forest, where all but one of them regularly run on different routes in the same forest, but the odd one out is running there for the first time, if the newcomer gets a very close fourth place, then it's not unreasonable to think that they're probably the best runner of all of them. ⁴

Finally -- I've been talking about the models as if they were all the same size (apart from the OpenAI medium weights), but there's one extra important difference with OpenAI's: they use weight tying, while mine don't. That means that their small model is significantly smaller than mine are -- closer to 124M parameters ⁵ than the 163M that mine have. Weight tying made a lot of sense back in 2019, when VRAM was even harder to come by than it is now, but when I looked into it, it made loss much worse.

Carrying on with the trail runner metaphor, the OpenAI small weights not only don't know the forest well -- they're also (scratches head looking for a good analogy) physically weaker. Their performance is even more impressive.

So: OpenAI weights good, mine (relatively) bad. Their smaller model, being tested for loss on a less-familiar-looking test set, does better -- or at least only a tiny bit worse than -- my larger ones. And when instruction fine-tuned, it converges on a better-performing result, and does it faster.

What could the difference be?

Probably not dropout...?

One thing that occurred to me while I was pondering this was dropout. If, when training, we randomly ignore 10% of the activations, perhaps we get a more generalisable model? Most of what I've read recently about dropout has focused on how it helps to avoid memorisation in multi-epoch training runs, but when I first read about it I felt that the generalisation argument was also strong. If true, that might explain the lower number of epochs needed for instruction fine-tuning, and maybe even the performance.

Luckily, though, I'd accidentally done an experiment that suggested that this wasn't the difference. As one of my JAX training runs, I'd trained one with dropout, and as you can see if you look back at the results table above, it was one of the worst performers in the IFT test.

One result is not a proof, but the fact that this one was bad compared to the others meant that I felt it was solid enough to ignore dropout as a possibility, at least initially.

Data quality?

Again, it's kind of annoying that WebText was never published. All we know about it is from the paper, where they say:

[W]e created a new web scrape which emphasizes document quality. To do this we only scraped web pages which have been curated/filtered by humans. Manually filtering a full web scrape would be exceptionally expensive so as a starting point, we scraped all outbound links from Reddit, a social media platform, which received at least 3 karma. This can be thought of as a heuristic indicator for whether other users found the link interesting, educational, or just funny.

How does this compare with FineWeb in terms of quality? Without seeing the results, we really can't say. FineWeb is more of a general web-scraping corpus, without the "filtering" provided by those Reddit upvotes. On the other hand, we're talking about Reddit here, and all kinds of awful junk gets upvoted, so I'm not confident that it's that strong a quality signal. And FineWeb has been curated to some degree.

That said, I had previously noticed that the models I had trained on the FineWeb-Edu dataset ("the most educational web pages" from FineWeb) punched above their weight in the IFT test. They got pretty poor loss on the test set, but didn't do too badly in IFT.

In retrospect, it's actually not that surprising that they did badly on the test loss eval; they were trained on a curated dataset of "good" stuff, and then were being evaluated against whatever less-curated stuff appears in the test set -- which you'll remember came from FineWeb. If you train a model on Jane Austen and then evaluate against Chuck Tingle, then you're not going to get amazing results.

But then, there's the excellent performance of the OpenAI weights on that same FineWeb test set. Even if WebText was higher quality than FineWeb, that clearly wasn't the whole story.

So data quality felt like it might be part of it -- but it wasn't where I wanted to start.

Overtraining?

This, finally, is where I landed for a first experiment. I'd been training my models on 20 tokens per parameter -- the Chinchilla-optimal number. But that heuristic dates to 2022, and GPT-2 was trained in 2019. While we don't know for sure, from what I've managed to dig up it looks like it was probably trained on much more data -- and probably over multiple epochs too.

Training on more tokens than the Chinchilla number is called overtraining (not to be confused with overfitting). So: what happens if we overtrain the model? Will we get the loss down? And if so, will we start approaching GPT-2 small's results on the IFT eval?

Stay tuned :-)

Here's a link to the next post in this series.

Modern production LLMs tend to use reinforcement learning as a big part of the fine-tuning phase, and that works quite differently. ↩
When I started reading up on this, I was hoping that it would be a manifold, largely because I like the word. And it is! A hypersurface is a specific kind of manifold. :-D ↩
If you like fantasy, there are some neat parallels there; the different overlaid worlds the characters move between in Philip Pullman's His Dark Materials series (aka The Golden Compass), China Miéville's Un Lun Dun, V. E. Schwab's Shades of Magic, and Charlie Stross's The Merchant Princes series come to mind. Even the Upside Down in Stranger Things works, though then you have to decide which of your loss landscapes is the evil one, and that muddies the analogy annoyingly. ↩
Sharp-eyed readers might have noticed that there are 16 models in the table above. I'm excluding the OpenAI medium weights. If you want to extend the analogy, that model is an Olympic runner who happens to have stopped by for the race... ↩
The paper says 117M, but the released weights are larger. I can successfully load them into my code with a model config that has 124,439,808 parameters. ↩

Why do OpenAI's GPT-2 weights beat mine? Part two: the bugfix

Thu, 30 Jul 2026 19:00:00 +0000

I'm digging into why my GPT-2 style models score worse on an instruction-following eval than OpenAI's original weights; I gave the details in this post.

While I was writing up the results of my first experiment into possible causes, I ran the post past ChatGPT -- I always use an "editorial board" of AIs to check my posts for flow, style, and any technical errors (though all writing is always mine). It took a look at the eval code that I was running, and highlighted a bug.

Luckily, it doesn't change the important results -- OpenAI's models continue to be better than mine at instruction-following. But it was enough to change the baseline numbers, re-ordering how well my own models did. So I fixed it and regenerated the baseline so that future experiments are based on solid ground.

The bug

The eval takes a model, and trains it over multiple epochs on a split of a subset of the Alpaca instruction-following dataset. At the end of each epoch, it evaluates the resulting model against a held-back validation split; if the eval loss starts rising, it bails out. Finally, it runs a test split of the dataset through the resulting model, and saves the result.

Once I've run it for a bunch of different models, I use an LLM-as-a-judge script to get GPT 5.5 to score results -- for each question-answering result, it sees all of the responses for all of the models in the same prompt, shuffled in order each time, to try to make it judge models against each other as consistently as possible.

Now, the idea was that the generation of the test split answers would use the model from the epoch prior to the rising-loss one. So I had code like this:

    for epoch in range(100):
        model.train()
        for input_batch, target_batch in tqdm(train_loader, desc=f"Epoch {epoch}"):
            optimizer.zero_grad()
            loss = calc_loss_batch(
                input_batch, target_batch, model, device
            )
            loss.backward()
            optimizer.step()

        model.eval()
        val_loss = calc_loss_loader(val_loader, model, device, eval_iter)

        if last_val_loss is None or last_val_loss > val_loss:
            last_val_loss = val_loss
            last_params = model.state_dict()
            print("Val loss still decreasing, continuing")
        else:
            print("Val loss rising, bailing out")
            break

    model.load_state_dict(last_params)

Each time the validation loss went down, I wanted to store the model's parameters in last_params so that they could be restored later in that last line.

If you look closely, the error is pretty obvious. model.state_dict() does not return a copy of the model's parameters -- it produces a dictionary containing references to the parameters inside the model. So although I was trying to stash away the parameters in last_params each time loss went down, so that we had a copy of the ones that were live in the epoch prior to the rising-loss one, what I was actually doing was just pointlessly saving a reference to the "live" params. The call to model.load_state_dict(last_params) was essentially a no-op.

The solution was simple enough:

from copy import deepcopy
...
        last_params = deepcopy(model.state_dict())

That was enough to make the code do what it was meant to do.

While I was there, I also noticed that the evaluation code was only using the first five batches in my eval dataset. This code was originally adapted from an eval in "Build a Large Language Model (from Scratch)", where the eval was run much more frequently and had to be super-fast. Because my own code ran it more rarely, it made more sense to use all of it. Given that I was going to completely re-run the script to generate the test responses, I figured that I might as well fix that at the same time.

A new baseline

I re-ran the fixed script on all of the models I'm comparing, and ran that past the GPT 5.5 judge; here's what I got:

	Test loss	Old IFT epochs	Old IFT score	Old IFT rank	New IFT epochs	New IFT score	New IFT rank
OpenAI weights: medium	3.231442	2	41.62	1	2	42.41	1
JAX, with MHA bias, no dropout	3.418784	4	19.25	4	4	18.12	5
JAX, no MHA bias, no dropout	3.420089	3	14.66	11	5	20.72	3
JAX, no MHA bias, with dropout	3.476802	4	12.94	15	7	16.98	8
OpenAI weights: small	3.499677	2	26.73	2	2	26.11	2
`1xrtx3090-stacked-interventions`	3.538161	4	17.79	6	4	13.68	11
`8xa100m40-stacked-interventions-1`	3.577761	4	10.29	16	4	9.81	15
Cloud FineWeb, 8x A100 40 GiB	3.673623	7	20.71	3	6	19.45	4
`1xrtx3090-baseline`	3.683835	6	15.11	9	6	13.39	12
`8xa100m40-baseline`	3.691526	4	14.74	10	4	14.64	9
Cloud FineWeb, 8x H100 80 GiB	3.724507	4	13.25	14	4	13.84	10
Cloud FineWeb, 8x A100 80 GiB	3.729900	4	14.50	12	4	11.30	14
Cloud FineWeb, 8x B200 160 GiB	3.771478	4	16.03	8	4	11.82	13
Local FineWeb train	3.943522	7	13.73	13	7	8.93	16
Local FineWeb-Edu extended train	4.134991	7	16.70	7	7	17.10	7
Local FineWeb-Edu train	4.166892	7	18.68	5	7	17.28	6

Let's dig into those numbers.

Training epochs

Let's look into the number of training epochs first; I've highlighted the models for which it changed.

For the "Cloud FineWeb, 8x A100 40 GiB" model, I think that was a result of the change to the number of validation samples.

For the two JAX models, it's a bit more of a mystery. There may be some part of the same "more eval data" variation there, but I have a suspicion that there's something more, related to dropout.

My methodology to date has been that the IFT training runs should use the same dropout setting as the original base model training run (I'll come back to this in a later post). However, due to differences between the JAX training script and the evaluation code (which is PyTorch), matching the dropout rate when evaluating those models is a bit fiddly and error-prone.

I am 100% sure that I got it right this time around -- I've checked and double-checked the specific commands I ran. But I think -- from what I remember doing -- that I might have messed it up the previous time. That's not certain, though -- just a suspicion based on half-memories of some commands I ran several weeks ago, which never wound up in my .bash_history (too many terminals open at once).

The ranking

For the IFT scores, remember that they're not strongly comparable between runs. Let's imagine that the LLM judge is given the following answers to the question "Who was the author of Pride and Prejudice?"

"Jane Austen"
"The author of 'Pride and Prejudice' was Jane Austen"
"The author of 'Pride and Prejudice' was Sarah Palin"
"The author of 'Pride and Prejudice' was 'Pride and Prejudice'"

In some cases it might treat the first two as being 100/100, in others it might give the second 95/100 for being too wordy. Likewise, in some cases it might rank the last two as 0/100 for being wrong, in others it might give the "Sarah Palin" one 5/100 for at least being the name of a person rather than complete nonsense.

Now, we always ask the LLM to judge all of the models' answers for a given question in one go, so at least we can be sure that it will be consistent for a given prompt about a given question. But what we can't keep consistent is which way it leans between different runs, or different questions within the same run. Sometimes it might be feeling "generous" and give the Sarah Palin answer a bit of grace, other times it might be harsher.

So there's a significant amount of noise there; my rule of thumb is that a variation of a point or two is within that noise, so OpenAI medium going from 41.62 to 42.41 is pretty much meaningless, and likewise "JAX, with MHA bias, no dropout"'s move from 19.25 to 18.12.

So what's important is the relative ranking -- which is first, which is second, and so on. Naturally, you have to allow for the fact that -- for example -- if one model goes from position 11 to position 3, like "JAX, no MHA bias, no dropout" did, then the previous number 3 will have to become number 4, 4 will become 5, and so on. You can see that happening in the results table. Obviously, as other models rise and fall, that has further knock-on effects.

Anyway, with all of those caveats, the good news is that my original mystery remains. The OpenAI models were still doing noticeably better than my own ones. GPT-2 medium continued to lead the pack (unsurprisingly, given that it's a bigger model), and GPT-2 small was still in second place. If that had changed, it would have made a rather disappointing end to this series: "mystery explained, it was a bug in the eval :-("

But now let's look at my models.

Firstly, it looked like the "no MHA bias" models might have benefited from the extra training -- or from having their dropout settings corrected. They rose from positions 11 and 15 to 3 and 8 respectively -- a huge swing for "JAX, no MHA bias, no dropout", and a solid improvement for "JAX, no MHA bias, with dropout".

Most of the other changes in relative rankings can be explained by those two models having been promoted, but there are some other changes. In particular, three models dropped significantly in score (and, as a result, ranking):

1xrtx3090-stacked-interventions
"Cloud FineWeb, 8x B200 160 GiB"
"Local FineWeb train"

My suspicion -- a very weakly-held hypothesis, but an interesting one -- is that those models had previously been benefiting from the bug.

Remember that the signal we're using to stop training is that the validation loss starts rising. We're using that as a proxy for overfitting, which in turn we're using as a proxy for "this model has had as much training on this data as it needs for the eval".

But there's no guarantee that the connection is there. Perhaps the models weren't overfitting and if we'd waited for another epoch or two, the validation loss might have started falling again. Or maybe some amount of overfitting would be beneficial for this eval?

Conclusion

There's probably a near-infinite amount of digging in that I could potentially do here. But I think it's best to stop. The bugfix was important because it meant that the eval was now doing what I thought it was doing.

Importantly, it doesn't change the puzzling fact that my models were worse at this eval than OpenAI's, which is what I'm trying to untangle. And it means that I can now lean more confidently on the baseline numbers.

So now it's time to actually start changing things to see if I can close the gap!

Here's a link to the next post in this series: does overtraining help?.

Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining

Fri, 31 Jul 2026 01:15:00 +0000

The GPT-2-style models that I've been training work really well, and I've even managed to train some that perform better than the original OpenAI small model in terms of cross entropy loss on a test set. But as I wrote previously, there's a mystery: why do they perform worse on my instruction fine-tuning evaluation?

I had various theories about why that might be, and to me, the most plausible-seeming of them was the amount of data they were trained with. As best I can find out, OpenAI's models were, by modern standards, trained on much more data than they should have been, while I'd used the theoretically optimal amount of training data.

To put it in other words, OpenAI's models were overtrained. If I deliberately overtrained my own models, could I match their performance? This post is a write-up of what happened, but so as not to bury the lede -- it didn't seem to help much, if at all.

Let's see why.

Overtraining

Let's start by getting a nice crisp definition of overtraining.

It's important not to confuse overtraining with overfitting. Overfitting is where you train a model so that instead of learning a general rule about the data it's seeing, it learns something very specific to the training data -- for example, this:

...rather than this:

Overfitting is pretty much always a bad thing.

Overtraining, by contrast, is more of a judgement call. For LLMs, it's generally used as a shorthand for "training for more than the Chinchilla-optimal number of tokens". The Chinchilla paper makes a very specific case: if you train a model for roughly 20 times as many tokens as it has parameters, then you'll have as good a model as you can get for that budget in terms of compute. They were arguing against contemporaneous experiments where people were doing things like doubling the number of parameters but training on the same amount of data.

If you overtrain, it means that you're training on more than 20 tokens per parameter. The Chinchilla argument is that instead of doing that, you should scale up the number of parameters and the number of tokens equally to keep the 20x ratio. Because the amount of compute used scales pretty much linearly with both tokens and parameters, you'll spend the same amount and you'll get a better model that way.

Based on that heuristic, if you double your compute budget, then you should scale up your parameter count by $\sqrt{2}$ and your training token count by the same amount, and by doing that you'll get a better result than you would if you'd naively just doubled the parameters or the tokens, for the same amount of compute time spent.

But overtraining is not always a bad thing. Keeping things Chinchilla-optimal means that you have to keep scaling up the model as you scale up the compute budget, and often you can't do that -- for example, let's imagine you're training a model that's meant to run on mobile phones. You have a hard limit on the number of parameters: what will fit in the target devices' RAM.

And importantly, in general you will still get a better model by overtraining -- just not as much better as you would have done if you had been able to scale up the model as well as the training tokens.

Now, as always, we don't know enough about the original GPT-2 training runs to be sure as to whether they were overtrained, and if so, by how much. But one thing that we do know is that GPT-2 was trained in 2019, three years before the Chinchilla paper came out, so they definitely didn't use it as a heuristic!

One thing that they do say in the GPT-2 paper is that their dataset, WebText, is "a total of 40 GB of text". Assuming 4 bytes per token (a good rule of thumb for the GPT-2 tokeniser), that's 10B tokens. If the small model, which had 124M parameters, was trained on all of those, then it definitely was overtrained; 124M times 20 is about 2.5B. ¹

On top of that, there's the question of epochs. In my training runs so far, I've been training through a Chinchilla-optimal 3.2B unique tokens. But being able to easily get your hands on that much training data is a relatively new thing, as you can see from the fact that the GPT-2 authors decided to document how they got theirs in the paper. So back then, pre-Chinchilla, people tended to train for multiple epochs so that they could make better use of limited data.

Now, when I first started training my models, I tried to dig up some details on the GPT-2 training run beyond what was there in the paper. I found this report, and using data there, I calculated that it looked like OpenAI had trained for about 42 epochs over WebText. The report's author came up with 60 epochs as an equivalent-sized training run for their own dataset. ²

Again, those numbers are shaky; we don't have the real data. But I think it's not crazy to say that the OpenAI models were probably trained on more data than mine, and probably over more than one epoch. They were, in the Chinchilla sense, overtrained. That makes perfect sense for the time, given that it was three years before the Chinchilla paper!

But that opened up a couple of interesting experiments that I could try.

The experiments

Obviously I didn't want to train a model over 10B tokens (not least because I'd need to download a larger sample of FineWeb). And I certainly didn't want to do 42 epochs of training, given that one epoch over 3.2B tokens took almost two days, even with my relatively fast JAX code.

But it seemed plausible that I'd be able to get at least some signal if I trained longer. I decided to use poppy, my new dedicated training box to train two new models:

Firstly, I'd train one on 6.4B tokens from my FineWeb dataset -- the original 3.2B that I had been training on to date, and then on whatever 3.2B came next.
Secondly, I'd train another one on the same 3.2B tokens as usual, but I'd do two epochs.

I expected each training run to take a bit less than four days on poppy. Once they were done, I would be able to evaluate the models, both with the test loss eval and the IFT one.

In an unusual-for-me fit of scientific good practice, I decided to write down what I expected to see up-front. I felt that:

Both models would get better results on the test loss eval than my existing ones (90% probability). The one trained on more tokens would be better than the one trained for two epochs on the same tokens (70%).
Both models would score better than my existing models in the IFT test (70%) but would still be worse than GPT-2 small (90%).

It was time to find out!

The extended-train model

I kicked off the extended train, over 6.4B unique tokens, using my JAX code because it's somewhat faster than the PyTorch version. It crashed after about 40 hours, with this error:

jax.errors.JaxRuntimeError: INTERNAL: CUDA error: Failed to end stream capture:
CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error
during capture [executable_name='jit_train_step']

There were no obvious issues with the machine -- the CPU temperature had been hovering at around 60°C, and GPU at around 70°C, as had been typical in training runs on poppy in the past. There was nothing in journalctl or dmesg that looked suspicious.

On Nvidia's documentation site I found a mention of that error message, but the page in question was all about PyTorch; I couldn't find anything relevant to JAX.

For now, I decided to chalk it up to some bug somewhere in the training stack, and only dig in if it happened again. So I kicked the training run off again from the latest checkpoint to see what happened, and 38 hours later, it completed:

Training complete in 140,667.401 seconds
2026-07-24 03:36:02.148127 Tokens seen: 3,191,341,056
2026-07-24 03:36:02.148133 Throughput: 22,687 tokens/second
2026-07-24 03:36:02.148155 Final train loss: 2.979
2026-07-24 03:36:02.148163 Done

Note that the numbers there -- tokens seen, time taken, and so on -- are for the portion of the training run after the restart. My checkpoint recovery code doesn't carry those over. The final train loss is just the loss over the period from the penultimate checkpoint to the end of the run -- there must have been some "easy" data there :-)

The training loss chart looked like this:

As you can see, the latest checkpoint wasn't the "best" one, so I pulled both latest and best down to perry, my workstation, for investigation.

Firstly, I ran them both through my JAX smoke test, which asks them to complete "Every effort moves you" with 20 more tokens (using greedy sampling):

giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run test_generation.py runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/best/model.safetensors
Every effort moves you forward.
The best way to get started is to start with a simple, easy to use online
giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run test_generation.py runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/model.safetensors
Every effort moves you forward.
- You are a leader.
- You are a leader.
- You are

Very inspirational. But coherent, and that's what matters.

Now, the bulk of my evals use PyTorch rather than JAX -- I've been sticking to that for consistency -- so I converted the model Safetensors files so that they had the right structure to work with that:

giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run convert_model_to_pytorch.py runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/best/model.safetensors runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/best/pytorch-model.safetensors
giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run convert_model_to_pytorch.py runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/model.safetensors runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/pytorch-model.safetensors

...and did the equivalent smoke test on that side (which uses non-greedy decoding with a temperature of 1):

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_smoke.py runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/best/pytorch-model.safetensors
Every effort moves you along to try the recipes for each and every one of them and try it all over again. You
giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_smoke.py runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/pytorch-model.safetensors
Every effort moves you to a conclusion.
If the answer for each of our goals is to get to where we�

I think that the Unicode junk at the end of the second is probably the first half of a two-token apostrophe or something along those lines. Anyway, those looked good.

Next, it was time to run the first eval: how would the models perform on the held-back test set?

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/best/pytorch-model.safetensors
Fetching 4 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 593.57it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [04:54<00:00, 10.88it/s]
Loss against our test dataset: 3.325897
giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/pytorch-model.safetensors
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 1508.74it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [04:54<00:00, 10.86it/s]
Loss against our test dataset: 3.324953

So, the "latest" checkpoint was better than the "best" one.

This is a problem with the way I define "best" in my training script. The original version of that script used a validation set to evaluate the model before every checkpoint; "best" at that point meant that the model was the best performer on that validation set. Later on, I decided to play a little loose with my training runs and pull out the validation -- this seemed reasonably safe because I was doing single-epoch training, so what I saw as the primary benefit of regular validation during training -- detecting overfitting by looking for rising validation loss -- was not so important. It's hard for a model to overfit on a single epoch training run.

However, doing that meant that "best" had to be changed to mean "best-performing on the training data". And that's actually a really bad metric, because the training data is different for each measurement, so they're not really comparable.

So, for this model, I decided that I'd discard the "best" checkpoint, and use the "latest" one.

But all of that aside, those numbers were pretty impressive! Not only did it beat my best Chinchilla-optimal model to date, which had got 3.418784 loss on the same eval, and the original OpenAI small model with a loss of 3.499677, but it was actually getting quite close to the OpenAI medium's loss of 3.231442. ³

Promising!

Update: The model is now on Hugging Face as gpjt/jax-with-mha-bias-no-dropout-extended.

So now it was time for the two-epoch training run.

The two-epoch model

Adding support for handling multiple epochs to the code was very simple, so having done that, I kicked it off.

Just over three days later:

Training complete in 282,366.091 seconds
2026-07-27 20:14:26.124497 Tokens seen: 6,520,504,320
2026-07-27 20:14:26.124503 Throughput: 23,092 tokens/second
2026-07-27 20:14:26.124518 Final train loss: 3.547
2026-07-27 20:14:26.124524 Done

This time the numbers are for the full training run -- there were no odd CUDA issues, and it ran straight through. Again, the final train loss is just the loss over the period from the penultimate checkpoint to the end of the run. For runs over the same dataset, it can actually be a useful quick-and-dirty way to compare results before running the full evals, but in this case it's obviously not comparable with the last run's result there, because we're talking about loss on different data.

Anyway, once again, "best" and "latest" were different checkpoints, as you can see from the loss chart:

...so I pulled them both to perry, did the first smoke test:

giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run test_generation.py runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/best/model.safetensors
Every effort moves you forward, and you’ll be able to see the results.
The best way to get
giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run test_generation.py runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/model.safetensors
Every effort moves you forward, and you’re not alone.
The most important thing is to be able to

...converted them to PyTorch:

giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run convert_model_to_pytorch.py runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/best/model.safetensors runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/best/pytorch-model.safetensors
giles@perry:~/Dev/jax-gpt2-from-scratch (main)$ uv run convert_model_to_pytorch.py runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/model.safetensors runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/pytorch-model.safetensors

Did the PyTorch smoke test:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_smoke.py runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/best/pytorch-model.safetensors
Every effort moves you to take on a larger responsibility.
We’re here to help you.
We�
giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_smoke.py runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/pytorch-model.safetensors
Every effort moves you closer, and you become a better person as the year progresses.”
There’s

So that all looked good (despite the Unicode junk), and it was time to work out the loss:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/best/pytorch-model.safetensors
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3507.68it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [04:55<00:00, 10.83it/s]
Loss against our test dataset: 3.332157
giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ runs/1xrtx3090-baseline/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/pytorch-model.safetensors
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 1760.09it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [04:53<00:00, 10.90it/s]
Loss against our test dataset: 3.326482

That was pleasingly in line with my predictions:

Both models would get better results on the test loss eval than my existing ones (90% probability). The one trained on more tokens would be better than the one trained for two epochs on the same tokens (70%).

...though of course the difference between the 3.326482 that this model got and the 3.324953 that the one-long-epoch one got was tiny and probably in the noise. I decided to count it as a win, anyway :-)

Update: The model is now on Hugging Face as gpjt/jax-with-mha-bias-no-dropout-2-epoch.

So now it was time for the big test: how would they do at instruction-following?

At last, the IFT test!

There are two phases to getting numbers for this test: firstly, I run the script that does the fine-tuning, and then generates the model's completions for the test set using the fine-tuned model. These are written to a JSON file for later use by the LLM-as-a-judge script. That script is the one I fixed in my previous post.

I ran it for the extended train (one long epoch) model, taking care to make sure that the config it was passed matched the model's original training run in not having dropout enabled:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run ift_generate_test_responses.py jax-mha-bias-extended-train runs/1xrtx3090-stacked-interventions/model.json ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-extended/checkpoints/latest/pytorch-model.safetensors
Epoch 0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.23it/s]
Val loss still decreasing, continuing
Epoch 1: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.54it/s]
Val loss still decreasing, continuing
Epoch 2: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.54it/s]
Val loss still decreasing, continuing
Epoch 3: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.58it/s]
Val loss rising, bailing out
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 110/110 [00:08<00:00, 13.40it/s]

Next, I ran it for the two-epoch model:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run ift_generate_test_responses.py jax-mha-bias-2-epoch runs/1xrtx3090-stacked-interventions/model.json ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-2-epoch/checkpoints/latest/pytorch-model.safetensors
Epoch 0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.13it/s]
Val loss still decreasing, continuing
Epoch 1: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.22it/s]
Val loss still decreasing, continuing
Epoch 2: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.34it/s]
Val loss still decreasing, continuing
Epoch 3: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.48it/s]
Val loss still decreasing, continuing
Epoch 4: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 116/116 [00:06<00:00, 18.32it/s]
Val loss rising, bailing out
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 110/110 [00:10<00:00, 10.06it/s]

With that done, it was time to run the LLM as a judge script. That gave us results, which I've put into the table below. For each model, I've shown the number of epochs of training the model needed before its validation loss started rising. For the pre-existing models, I've shown the score that they got in my baseline evaluation at the end of my last post, and their ranking in that eval. Then for all models, there's the score from this new run, and the new ranking. The two new models are in bold.

	Test loss	IFT epochs	Old IFT score	Old IFT rank	New IFT score	New IFT rank
OpenAI weights: medium	3.231442	2	42.41	1	42.91	1
JAX, overtrained one long epoch	3.324953	3			18.75	5
JAX, overtrained two normal epochs	3.326482	4			18.45	6
JAX, with MHA bias, no dropout	3.418784	4	18.12	5	17.53	8
JAX, no MHA bias, no dropout	3.420089	5	20.72	3	20.25	3
JAX, no MHA bias, with dropout	3.476802	7	16.98	8	16.73	10
OpenAI weights: small	3.499677	2	26.11	2	25.66	2
`1xrtx3090-stacked-interventions`	3.538161	4	13.68	11	13.35	14
`8xa100m40-stacked-interventions-1`	3.577761	4	9.81	15	10.33	17
Cloud FineWeb, 8x A100 40 GiB	3.673623	6	19.45	4	19.76	4
`1xrtx3090-baseline`	3.683835	6	13.39	12	13.74	13
`8xa100m40-baseline`	3.691526	4	14.64	9	14.11	11
Cloud FineWeb, 8x H100 80 GiB	3.724507	4	13.84	10	13.84	12
Cloud FineWeb, 8x A100 80 GiB	3.729900	4	11.30	14	11.13	16
Cloud FineWeb, 8x B200 160 GiB	3.771478	4	11.82	13	11.22	15
Local FineWeb train	3.943522	7	8.93	16	8.73	18
Local FineWeb-Edu extended train	4.134991	7	17.10	7	17.42	9
Local FineWeb-Edu train	4.166892	7	17.28	6	17.61	7

A note on the numbers: in my previous post, I mentioned that different LLM judge runs differ due to randomness in how "strict" the judge model is. You can see that showing up here -- most of the old/new IFT scores are pretty close, within a point or so, but they do differ, some rising and some falling. This is with exactly the same set of responses for each model going into the judging program -- I used the same JSON files for this table as I did for the previous one for the models that were in there -- the only difference was that the new JSON files for the new models were added.

As you can see, the new models scored better than the "JAX, with MHA bias, no dropout" one, which is the most similar: it's the same model config, just trained for the Chinchilla-optimal number of tokens. The one long epoch model got a score that was 1.22 higher than that one, and the two-epoch one scored 0.92 higher.

However, my normal rule of thumb for comparing models in these evals is that differences of less than a point or two are probably in the noise.

I did a second run of the LLM judge script -- it takes 20 minutes to run and costs a couple of dollars each time, so I don't like to run it all that often -- and this time around (just looking at the JAX numbers) things were a bit closer:

	IFT score
JAX, overtrained one long epoch	18.45
JAX, overtrained two normal epochs	19.45
JAX, with MHA bias, no dropout	18.30
JAX, no MHA bias, no dropout	21.04
JAX, no MHA bias, with dropout	17.33

You can also see that the two new models have swapped places.

So I think that the principled approach here is to say that the improvement is almost certainly in the noise. Perhaps if I did a very large number of runs of the judge I'd get something more solid -- but what I'm hoping for is something a little more unambiguous; some change that really moves the needle in an obvious way, clearly outside the noise.

That means that my second pre-registered prediction:

Both models would score better than my existing models in the IFT test (70%) but would still be worse than GPT-2 small (90%).

...was half wrong. I got the "worse than GPT-2 small" bit right, at least. But while the models did look a bit better than the most similar pre-existing model, (a) the difference was too small for me to be confident in it, and (b) they were still worse than "JAX, no MHA bias, no dropout" and "Cloud FineWeb, 8x A100 40 GiB".

What can we take away from this?

Conclusion

The hypothesis that I was trying to test was whether it was simply overtraining that made the OpenAI weights better at this IFT evaluation than mine. Frustratingly, I can't say that the hypothesis was false.

Perhaps there is an improvement gained by overtraining -- and the apparent gains which, on this experiment, appeared to be in the noise, would have been consolidated if I'd trained for even longer -- perhaps the 42 epochs on 10B tokens I suspect that the original weights were trained on?

But equally, perhaps there is no benefit, and further training would have left the models exactly where they were in the ranking.

Intuitively, you'd think that further training of a model would make it better at answering questions, at least up until the point that its parameters were "saturated" and could not absorb new information without forgetting something else. After all, a model that's never seen "Jane Austen wrote 'Pride and Prejudice'" will never be able to successfully answer when it's asked who the book's author was.

But where that saturation point might be -- and indeed how much training you'd need to do to get there -- is not obvious.

It's an annoying place to finish this experiment, but I guess at least an inconclusive result is better than never having run it at all. And it was at least good to see the test loss improvement that I expected.

But given the opportunity cost of tying up poppy in four-day training runs, I think I'll look into other possibilities next. As I was running this experiment, something came up -- and I'll post about that soon. Luckily, this time it won't involve training more base models...

Here's a link to the next post in this series: digging into dropout.

It's also worth noting that by the same, um, token, the extra-large model, with 1,542M parameters, was undertrained because the Chinchilla-optimal number of tokens would have been about 31B. Though that said, see later regarding epochs. ↩
You might wonder whether training over the same tokens repeatedly over multiple epochs "counts" for Chinchilla purposes. Is training on 1.6B tokens over two epochs the same as training on 3.2B tokens over one? "Scaling Data-Constrained Language Models" poked into that in 2023, and from the abstract, came to the conclusion that you could do up to four epochs over the same data without losing much value, but after that returns diminished. If that holds for GPT-2, and they really did train for 42 epochs, maybe they wasted a lot of time? I'll need to read that paper in full at some point. ↩
There is a table of results later on in this post where you'll be able to compare models easily. ↩

How I use AI on this blog

Fri, 31 Jul 2026 19:00:00 +0000

Inspired by this LessWrong post, I thought I'd write about how I use AI here. This is less in the interest of disclosure, more to provide a snapshot of what I'm doing right now so that I can revisit it in the future and see how it changes. And hey, maybe it'll be of interest to you, dear readers.

If I were to summarise my working philosophy in fewer than ten words, it would be: AIs identify problems and I fix them myself. With a very specific kind of exception (which I always flag), the text and code on this blog are human-generated. That's not a moral stand, but more a constraint imposed by what this blog is meant to be -- a place for me to learn in public.

Ideation and running experiments

Every post here is based on an idea I had, and work that I've done. For many posts -- for example, the large-scale coding projects like this one -- I'll have multiple chat sessions ongoing while I do the work, normally with either ChatGPT or Claude, or sometimes both. The amount of input they have varies, but because the value of these projects is in what I learn when I'm doing them, letting an AI do my thinking for me would make the whole thing pointless -- so I take steps to stop that from happening.

AIs are, of course, trained to be helpful, and will often explain things in their replies that I would have better learned on my own through experimentation. I'm generally pretty good at spotting when that happens before I've read more than a sentence or two, though, so I can skip reading that part, scroll straight down to the input field, and ask it to operate in more of a rubber duck mode.

With the most complicated projects, where each step has a hard dependency on having got the previous one just right, I do use AIs for code review. Let's say that I've built a model that I intend to extend. I'll test it myself (does the loss go down when training, is it generating plausible-looking results?), but if I want to be really cautious, then I'll run the code past an AI. I'll paste it into a chat session and tell the LLM what it's meant to do, and ask it to check if I've screwed up ¹. Again, though, I make it clear that I don't want it to make fixes -- just to point out any bugs.

Writing things up, and the editorial board

As things progress with a project, I keep detailed notes. When I'm done, I write them up without AI assistance, getting the post to a level where it might be a little messy in terms of how things are explained, but all of the important information is in there. I read it through and make sure I'm reasonably happy with it, and then it's time for what I've taken to calling the editorial board.

I paste the draft post into a fresh chat session with an LLM -- right now, this is normally Claude -- and ask for comments. It already has enough information in its memory of earlier conversations to know that what I'm looking for: places where I'm confidently wrong or other technical errors, places where my explanations are missing a step, or where I'm overexplaining things, conclusions that don't really follow from the results of an experiment, and that kind of issue. It tends to spot a few silly grammatical errors and typos at the same time. An important standing instruction is that I do not want it to rewrite anything. Just as with the code, it should tell me where there is a problem, and let me fix it.

We iterate on that for a while until we have something that we're both happy with, and then I feed it to the next LLM -- normally ChatGPT.

ChatGPT has a much more pernickety attitude than Claude does. I often use metaphors, and it will generally want me to replace them with mathematically rigorous prose. This is still very useful, though. Sometimes the metaphors aren't flagged as such well enough -- or even worse, there are times when the terms I've hit on for a metaphor happen to clash with a technical term, making what I've written misleading at best.

I don't always address all of the issues that ChatGPT raises, as otherwise every post would be a mess of hedges and overexplanations and what-have-you, but I like to get to a stage where I'm comfortable that I have a good understanding of specifically why I'm rejecting the remaining points it raises.

One other point where ChatGPT has helped a lot is that it's very diligent about checking supporting materials that I link to. When I was recently about to post an article on running an eval on a model, it followed the link to the training code and spotted a silly bug. It wasn't something that materially changed the eval's outcome, which is probably why I'd not noticed it, but it was something that was important to get right if I wanted later runs of the same evals to be solid. Definitely helpful.

With that done, I run it past a cast of other LLMs. The exact set varies over time; for the last few posts it has been (in this order) DeepSeek, Grok, GLM-5.2 and Kimi K3. I did use Gemini in the past, but over time it became less effective and just started complimenting me on the post and suggesting related topics to chat about, which was kind of pointless. I'll wait until the next release and then try it again.

Because the Claude and ChatGPT passes have generally got rid of anything particularly nasty, this second group of AIs often don't have much to add. However, occasionally they will spot something the others have missed, or have other suggestions, so it's worth spending the five minutes or so it takes to use them. It also helps to keep me up to date with what the other models out there are like. ²

When all of that's done, I run it past Claude one final time, tidy up any remaining issues, then publish it on a private staging site, and read through it carefully myself. The best time for that final readthrough is after dinner, ideally after a glass of wine; the goal is to smooth out the prose, and remove anything overly formal. To make it as close to being fun to read as I can manage.

Once I'm happy, I can promote it to the live site and hit the publish button.

That probably all sounds much more complicated than it actually is. A short post will normally go through all of that in half an hour -- less if I skip the full editorial board, which I sometimes do. The longer ones can take an hour or two, but given that they're normally the result of a week of work, on and off, in percentage terms it's not that much, and it's worth it for the polish.

So, there's no AI-generated text here, but I do lean on AIs to make it the best version I can of what I have to say. How about the code?

AI coding

Again, the goal of the projects I document here is to learn in public. If I'm learning some concept that is expressed in code, then I need to write that code. So that means that anything non-trivial will be something I've written by hand, with AI input limited to code review -- the same rule as I have for the text.

Of course, sometimes there are things I'd like to publish that I wouldn't learn anything by writing. Coding up matplotlib stuff to chart loss curves, or writing a fancy JavaScript visualiser to show what models' parameters are used for would teach me nothing. So for that kind of thing, I just let the AIs get on with it (and even then, for the parameter visualisation, I hacked a first version together in a spreadsheet to check my understanding, and then tested the visualiser against it). I do always mention in the text when a particular bit of code was AI-generated, though.

So, my rule is: if I would learn something by writing the code, I'll write it. If not, I'm happy to delegate to an AI.

But even then, I apply one restriction: if it's for the blog, I'll ask for the code in a chat session, rather than using a more agentic system like Claude Code or Codex. This is to add friction. If you're using an agent, it's easy for a task to grow, and what started as a throwaway idea can come to consume more and more time and cognitive space. Keeping it in the chat interface keeps things minimal -- or at least, that's how it works for me.

Does that mean that I'm against coding agents?

Yes to agents, but not on this blog -- yet

This blog is where I post about experiments I've done and what I've learned. At the moment, what I'm learning is all pretty low-level. How does an LLM work? What factors make it smarter or dumber? It's all pretty hands-on, and involves code that I need to understand.

I do other things apart from writing this blog, of course :-) And for that I'm keen on agentic tools; I have an OpenClaw agent to help me run my life generally, and use Codex and Claude Code for projects where I'm trying to achieve a specific goal, rather than trying to learn something. But by their very nature, those are not projects that will wind up here on the blog right now.

That might change in the future! When I feel that I have a solid, large enough foundation, perhaps I'll be running experiments that I want to write about, where it would make sense for AIs to handle the details, while I focus on the broader strokes.

But that time is not now, so right now, you can be sure that every word ³, and almost every line of code, was written by hand. Even if I do need the AIs to keep me on track and at least borderline coherent.

If you're wondering "why not use Claude Code or Codex", I get into why I avoid them for blog-related work towards the end of this post. ↩
Current thoughts:
- DeepSeek used to be close to Claude and ChatGPT, but has been left somewhat behind. I have been hearing rumours on X of an upcoming update, though.
- Grok used to have a propensity to try to turn my posts into clickbait. It would always want to rewrite things, and would suggest titles that were not a million miles away from "Ten things you never knew about RNNs -- number four will shock you!" Recent releases have been much better, though, and I might experiment with moving it further forward in the sequence.
- GLM-5.2 recently complimented me on my "science fiction" blogpost that mentioned ChatGPT 5.6 Sol. I probed it a bit on that and it said that because the publication date on the post was in 2026, it understood that I was writing near-future SF. From its perspective, the real date was sometime in late 2024. Surprising! I think the last time I saw that kind of behaviour from a model was, well, sometime in 2024...
- Kimi K3 is really quite impressive. I will be using it more. I particularly like the way it shows a fairly detailed chain of thought -- something it shares with DeepSeek and GLM-5.2, but there seems to be more depth there. It's a pity that Claude and ChatGPT only show summaries in the chat interface, though I understand their reasoning.
↩
ChatGPT, being wonderfully true to form, thinks I should clarify here that I'm not including quotes from models here (like the ones here) when I say that. ↩

A quick(ish) Chinchilla check

Fri, 07 Aug 2026 19:00:00 +0000

I recently overtrained a couple of GPT-2 style models, training them both on 40 tokens per parameter rather than the 20 per parameter that is generally regarded as "Chinchilla-optimal".

The normal heuristic is that instead of doing that, you should scale up the number of tokens and the number of parameters equally -- so I would have been better off scaling up the model by $\sqrt{2}$ and the token count by the same amount. By doing that, I should expect to get a better model in terms of loss on my held-back test set than I did with my 40-tokens-per-parameter models.

My training machine poppy wasn't doing anything, so I decided to give that a go. Would the Chinchilla rule-of-thumb hold up?

As you might expect, it did. But it was a surprisingly close-run thing, and could conceivably have been in the noise. Let's take a look.

The Chinchilla heuristic

If you already know all about the Chinchilla paper -- regular readers in particular must be sick and tired of it by now :-) -- then click here to skip this section.

In "Training Compute-Optimal Large Language Models", which is always called the Chinchilla paper after the name of the model they trained at the end, the authors tried to work out the optimal number of tokens to train an LLM on based on its number of parameters. In particular, they were pushing back on a trend they were seeing at the time, where people were making models ever-larger, but not increasing the amount of data they were training on.

The authors were all at Google DeepMind, and this was the kind of project that only a large lab could do: they trained "over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens". Their conclusion was "for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled".

They don't actually state an overall optimal number of tokens to train on in the paper, but in table 3 they provide an estimate of the optimal training FLOPs and tokens for models of various sizes, and it's approximately 20 tokens per parameter.

That number has become a heuristic, and people talk about a model as being trained for the Chinchilla-optimal number of tokens. Models that were trained on fewer tokens per parameter are referred to as "undertrained", and models that were trained on more as "overtrained".

It's worth noting that overtraining a model is not, in itself, a bad thing. If you have a model of a particular size and you continue training it past the Chinchilla-optimal number of tokens, it will -- in general -- get better. The point of the heuristic is that doing that is not the best way to spend whatever budget you have in terms of compute time. You'll get better results, as they say, by scaling the number of tokens and the number of parameters equally.

But let's say you're creating a model for specific target hardware -- say, a mobile device. You have a hard restriction on how large the model can be -- the device has only so much RAM to hold it. So it might make sense to overtrain to get a better model. ¹

But if you're not so limited in how many parameters you can use, then you should indeed scale the model up, and that's what I wanted to try. How would that work?

Scaling the model

A week or two back, I was investigating whether I could make my GPT-2 style models better at a specific instruction-following task by overtraining them. The details of that experiment aren't important here, but what it meant was that I had three GPT-2-style models, each of exactly the same size, roughly 163M parameters

A Chinchilla-optimal one, which I'll call jax-gpt2-chinchilla here.
One trained on twice the Chinchilla-optimal tokens, jax-gpt2-2x-chinchilla.
One trained on the Chinchilla-optimal tokens, with two epochs (so that it was trained for as long as #2): jax-gpt2-2-epoch-chinchilla

When I tested them against a held-back test set of sequences -- stuff that they'd never seen before -- they got results rather like you might expect:

	Test loss
`jax-gpt2-2x-chinchilla`	3.324953
`jax-gpt2-2-epoch-chinchilla`	3.326482
`jax-gpt2-chinchilla`	3.418784

A lower loss is better, and you can see that the longer-trained models were noticeably better than the Chinchilla-optimal one. The difference between them was tiny; they were trained starting with the same initial weights, and the training runs themselves were deterministic, but a difference of 0.05% in loss doesn't seem like it could be meaningful -- an extra batch for one or one fewer for the other could easily swap them around, you'd think.

Now, these models each had 163,009,536 parameters -- they were the small-size model from the GPT-2 paper, modified to not have QKV bias or weight-tying. jax-gpt2-chinchilla had been trained on 3,260,190,720 tokens (rounded up to fit into a round number of full batches), and the other two on 6,520,381,440 tokens each -- double the amount (rounded up too).

What I needed to do for my Chinchilla check was to try training a model that used the same amount of compute, scaling the parameters and the number of training tokens equally. Because training compute increases roughly linearly with both parameters and tokens, that would mean scaling both up by $\sqrt{2}$ , giving us:

163, 009, 536 * \sqrt{2} \approx 230, 530, 296 parameters

...and thus 4,610,605,920 tokens.

How to scale the model up?

In the GPT-2 paper, they train four models:

Name	Parameters ²	Layers	`d_emb`	MHA heads ³
small	124M	12	768	12
medium	345M	24	1024	16
large	762M	36	1280	20
xl	1542M	48	1600	25

I wanted to scale my own model up from 163M parameters to about 231M. Which of those numbers would I want to increase, and by how much?

The first thing that stands out is that the number of heads is always 1/64th of the number of embedding dimensions. So that sorted that one out. I just needed to adjust the number of layers, and the number of embedding dimensions, but ensure that the latter was a multiple of 64.

I decided to see if I could fit some kind of curve to the relationship between the number of parameters and the GPT-2 authors' choices. This was made a bit more complicated by one thing: they were using weight-tying, and I was not. That meant that they re-used the embedding matrix at the start of the LLM as an output head at the end -- which is why they had 38M fewer parameters. Embeddings and the output head make up a surprisingly large percentage of the parameters for small models like this -- about 47% without weight-tying, 23% with.

I couldn't work out a solid way to scale things up and wound up doing some rather messy hacking around in a spreadsheet. I came up with two proposed model sizes that were within a couple of percentage points of the right size:

Name	Layers	`d_emb`	MHA heads	Parameters	% diff
`slightly-larger`	15	896	14	235,621,120	+2.21%
`slightly-smaller`	14	896	14	225,978,368	-1.97%

Interestingly, I found that because d_emb could only change in increments/decrements of 64, it was a pretty coarse control -- my first attempt at making a slightly-smaller model changed it to the next step down, 832, but that led to a model that was 9.25% too small.

That was an interesting first lesson. I'd previously been thinking of the Chinchilla rule as being something like "don't double the tokens, just scale the model and the tokens equally". But that "just" was wrong. Scaling a model is hard -- even with just two dials to fiddle with, like in this case, it was tricky to get something right -- and I can't say for sure that my choices were the right ones.

Anyway, the next step was to double-check that these models would use the right amount of compute to train.

Training FLOPs

As I said earlier, the compute time scales roughly linearly with the number of parameters. Let's dig into that "roughly".

Different kinds of parameters take different amounts of FLOPs to train, and scale differently with things like the embedding dimensions, sequence length, and so on.

Now, for very large models, a lot of that comes out in the wash, but with tiny models like these where the embeddings make up such a large proportion of the parameters, it might matter.

Conveniently, in appendix F of the Chinchilla paper, they provide a set of formulae for estimating the number of training FLOPs for a normal dense LLM like these ones. I coded that up into a script that, given the JSON configuration files I was using for my models and training runs, would work out the number of FLOPs for a single epoch of training. It didn't take account of the fact that my real training runs round the number of tokens up so that we do a round number of full batches, but I felt that so long as the results weren't very close that wouldn't matter.

I got these results (multiplying the two-epoch numbers by two):

	Est. training FLOPs
`jax-gpt2-chinchilla`	3,544,967,596,946,227,200
`jax-gpt2-2x-chinchilla`	7,089,935,193,892,454,400
`jax-gpt2-2-epoch-chinchilla`	7,089,935,193,892,454,400
`slightly-larger`	7,419,664,885,127,577,600
`slightly-smaller`	6,804,429,215,367,168,000

The numbers were indeed different enough that I wasn't worried about the batch-rounding. And the good news was that slightly-larger and slightly-smaller would indeed use slightly more and slightly less compute to train than the overtrained models -- about 4.6% more and 4% less respectively. A true Chinchilla-equivalent model would lie somewhere between them.

It was time to train some models!

Training

I kicked off the run for the slightly-larger model first. Because it was bigger than the 163M models I'd been training, I couldn't fit such large batches into my VRAM; previously I'd been running with a batch size of 6, and now I could only fit in a batch of 4. Luckily, though, I was using gradient accumulation, so by bumping that up from 16 steps to 24 steps I could keep the same overall batch size and keep the training runs comparable.

Even despite that, the training run ran out of VRAM about 60 hours in -- I'm guessing due to VRAM fragmentation, as I did not have TF_GPU_ALLOCATOR set to cuda_malloc_async -- but I was able to restart from the most recent checkpoint and complete the run. After just less than four days total training time, it completed.

When it was done, I copied the last checkpoint ⁴ over to my dev box, perry, and ran my standard smoke test against it, asking it to complete "Every effort moves you" with 20 tokens, using greedy sampling. I got something reasonably coherent:

Every effort moves you.
I’m not sure what you’re thinking.
I’m

Next, I converted the safetensors file -- which had been saved by my JAX code -- into a format compatible with my PyTorch code, because that's what I use for evals. I ran another smoke test (this one with temperature 1):

Every effort moves you through the motions for your life, your soul, your body,
and your soul’s happiness

Very spiritual. Next, it was time to work out the loss on my held-back test set:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla/checkpoints/latest/pytorch-model.safetensors
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3485.09it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [07:11<00:00,  7.42it/s]
Loss against our test dataset: 3.280028

Well, it was certainly better than the 3.324953 that the best of the overtrained models got -- but only by a bit over 1% better. Interesting!

Update: The model is now on Hugging Face as gpjt/jax-with-mha-bias-larger-chinchilla-1.

I decided to train the second model, slightly-smaller. This one crashed mid-way through with an error that I've seen before:

jax.errors.JaxRuntimeError: INTERNAL: CUDA error: Failed to end stream capture: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error during capture [executable_name='jit_train_step']

I'm going to have to investigate that more in future, but for now, I just restarted from the checkpoint, and again after a bit less than four days, I had a model.

The JAX smoke test was solid:

Every effort moves you forward.
The best way to get started is to start with a free trial.
You can

...and so was the PyTorch one:

Every effort moves you forward in love with our products.
I love the way it’s easy to use.

Both quite commercial this time! It was time for the proper test loss eval:

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla-2/model.json ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla-2/checkpoints/latest/pytorch-model.safetensors
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3151.24it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [06:45<00:00,  7.90it/s]
Loss against our test dataset: 3.292937

So, slightly worse than the 3.280028 from the larger model, better than the 3.324953 from the best overtrained one.

Update: The model is now on Hugging Face as gpjt/jax-with-mha-bias-larger-chinchilla-2.

Time to put this all together.

Results

Here's an updated version of the table from the start of this post; I've added in the two new models, and the improvement they each had over jax-gpt2-2x-chinchilla in both absolute terms and as a percentage rounded to 3sf.

	Test loss	Improvement	Improvement %
`slightly-larger`	3.280028	0.044925	1.35%
`slightly-smaller`	3.292937	0.032016	0.962%
`jax-gpt2-2x-chinchilla`	3.324953	-	-
`jax-gpt2-2-epoch-chinchilla`	3.326482	-	-
`jax-gpt2-chinchilla`	3.418784	-	-

Now, unlike the overtrained models, prior to training these two new ones started with different initial weights to the jax-gpt2-chinchilla one -- after all, they had to, because they had more of them!

A while back, I did a bit of analysis of how random variation in weight initialisation can change the resulting test loss. It wasn't anything in-depth, but I trained three models with different explicit seeds set prior to the model initialisation, but with the same seed set before the training run started ⁵.

Those three models wound up with test losses of 3.681356, 3.673943, and 3.664345. Doing statistics with three data points is a bit flaky, but the cost of training models is so high that I'll leave the Proper Science to the likes of Google DeepMind and wing it :-)

Mean: ~3.673215
Sample variance: ~0.000073
Standard deviation (SD): ~0.008529

Now, piling statistical flakiness on statistical flakiness, we'll compare these. You'd normally expect about two thirds of results to be within one SD of the mean, 95.4% to be within two SDs, and 99.7% to be within three.

Three SDs on that (yes, different, I know) distribution is 0.025587. That's smaller than both of the improvements that our Chinchilla-optimal runs had over the overtrained ones.

So what does that tell us? Well, perhaps not much given the statistical flakiness. But I think it is useful directionally. It suggests that we might be able to take these results seriously as an improvement, and that Chinchilla held: scaling up the model and the number of tokens evenly did give us a better model than just scaling up the number of tokens. In particular, the fact that the loss for slightly-smaller was lower -- even though it had 4% less compute spent on it than the overtrained models -- was encouraging.

But it's certainly far from a slam-dunk. A larger test, training lots of overtrained models and lots of Chinchilla-optimal ones, all with different random seeds, would give actual real serious data.

Not worth it for me, and perhaps not for anyone.

Conclusion

I wanted to do a quick sanity check of the Chinchilla heuristic of 20 tokens per parameter. I came up with results that were certainly in line with it -- perfectly so in terms of the ordering of the models I trained. But the effect was small enough that I could imagine that it was in the noise, especially given the small numbers of models I'm able to train. I'll chalk it up as a tentative success.

In addition, I learned one useful thing: when talking about scaling up a model to more parameters, you actually have to think quite hard about where you want to put those parameters. I wound up doing a rough curve-fit to the models in the GPT-2 paper, but I have no idea if that was optimal. At some point I should try to dig up some research into optimising embedding dimensions, numbers of layers, and so on. But not now, as I've a bunch of other stuff I want to investigate first.

Anyway, I hope you found this experiment interesting, and as ever, comments and questions welcome below. Thanks for reading!

I'm less familiar with arguments for under-training -- that is, for fewer than 20 tokens per parameter. I've heard that these days, modern LLMs get a lot more reinforcement learning than they do pre-training, and perhaps that might mean that some very big ones are undertrained prior to RL? I'm uncertain. It's unlikely to be raw lack of data; even for those of us outside the big labs, FineWeb has 18.5T tokens. On its own, that would be enough to train a 0.925T-parameter model, and given that you can apparently do four epochs over the same data before you start getting diminishing returns, that takes us up to 3.7T. That's frontier-lab size, and I'm sure they have better datasets than FineWeb. ↩
Parameter counts are from the paper, apart from the "small" model, which is known to be wrong -- I used my own calculation, and the result is in line with what I've seen elsewhere. ↩
The paper doesn't mention the number of heads; these numbers are from "Build a Large Language Model (from Scratch)", and match up with the ones on this Hugging Face page. ↩
Regular readers might have noticed that I'm ignoring what I've been calling the "best" checkpoint. I've come to the conclusion that because for my training script, "best" means best in terms of training loss, and the training loss changes based on what training data the model has seen recently, it's actually not a very useful metric and just confuses things. At some point I'll probably re-introduce pre-checkpoint evals and use that for "best", which would be the right way to do it. ↩
At the time I was using dropout, so training runs were not deterministic without a known seed. ↩

Use the built-in GELU, don't roll your own!

Thu, 20 Aug 2026 02:00:00 +0000

Unsurprisingly, PyTorch's own built-in GELU function is faster than the hand-rolled one I've been using to date. But I was surprised at how much faster using it made things when training my models. I discovered this accidentally just now while working on something unrelated, but am logging the details here for anyone else that might find it useful.

The headline numbers: the same code, training the same model on the same data, ran at about:

21,000 tokens per second using the hand-rolled GELU from Sebastian Raschka's book "Build a Large Language Model (from Scratch)".
25,000 tokens per second using PyTorch's built-in GELU with no arguments.
25,000 tokens per second using the built-in GELU with approximate="tanh", which uses the same maths as Raschka's version under the hood.

That's a 20% increase in throughput for both of the built-in versions -- definitely nothing to be sneezed at.

And what is particularly interesting is that there aren't that many GELUs going on -- it's a GPT-2 small-style model, with 12 layers. So that's 12 GELUs handling tensors shaped (batch_size, seq_len, 4 * d_emb), which is (6, 1024, 3072) for my training setup. Given that the rest of the model is doing all of the normal full attention stuff for GPT-2, it's really surprising that the GELUs alone must have been taking up so much of the time. The throughput numbers mean that we must have been spending about 17% of our time on the extra overhead from the hand-rolled version, so that sets a lower bound for how much time the GELUs were taking up.

Back when I was doing the "interventions" part of my LLM from scratch series, training dozens of GPT-2 small-sized models in the cloud and on my local machines, to keep things simple I used the original model code from Raschka's book. That happens to have its own implementation of the GELU function -- you can see my copy here.

I'm not that sure why the hand-rolled version is in there -- he covers the maths, but the specific implementation isn't explained in that much depth, and it seems rather like boilerplate, just a "type this in and use it" kind of thing. By contrast, for example, while he does explain the maths behind cross-entropy loss in similar detail, we use the built-in function for it rather than coding it up ourselves.

When I switched to using JAX for my own from-scratch implementation, I decided to not bother porting the boilerplate, and just used JAX's own built-in version.

I was revisiting the PyTorch code -- I'm in the process of extending it with mixture-of-experts support, about which more in a later post -- and decided to switch from the hand-written GELU to the PyTorch one just to tidy things up a bit. I noticed something interesting -- my new MoE code suddenly seemed to speed up.

Was that a mirage? Or had I discovered part -- or even all -- of the reason why the JAX code was so much faster than the PyTorch code? With PyTorch, I was typically getting training speeds of about 21,000 tokens per second, while in JAX I was getting 24,000 tps or so. I'd been chalking that up to JAX's JIT compilation, but could it have been just a result of a random implementation choice I'd made?

I did three partial test training runs, letting each one run for 20 minutes to allow the training speed to settle down from any startup overhead. Firstly, with the old hand-coded GELU:

giles@poppy:~/Dev/ddp-base-model-from-scratch ((HEAD detached at 16dd249))$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3173.90it/s]
Starting rank 0 training at global step 0
  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=19,982]


Checkpoint

Continuing training
  1%|▊                                                                                                     | 257/33165 [20:07<42:53:57,  4.69s/it, loss=6.570, tps=20,920]

So it was getting 20,920 on average over those 257 global steps. That speed was in line with the original run of the configuration I was using.

Next, I introduced the built-in PyTorch GELU with no arguments:

-
-class GELU(nn.Module):
-
-    def forward(self, x):
-        return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2.0 / torch.pi)) * (x + 0.044715 * torch.pow(x, 3))))
-
-
-
 class FeedForward(nn.Module):

     def __init__(self, cfg):
         super().__init__()
         self.layers = nn.Sequential(
             nn.Linear(cfg["emb_dim"], cfg["emb_dim"] * 4),
-            GELU(),
+            nn.GELU(),
             nn.Linear(cfg["emb_dim"] * 4, cfg["emb_dim"])
         )

That does the full calculations for GELU, rather than using the tanh-based approximation that the hand-rolled code did. After 20 minutes, it looked like this:

giles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/
Fetching 4 files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 11222.22it/s]
Starting rank 0 training at global step 0
  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=23,953]


Checkpoint

Continuing training
  1%|▉                                                                                                     | 307/33165 [20:00<35:40:59,  3.91s/it, loss=6.331, tps=25,134]

So this time we were getting 25,134 tokens per second -- 20% faster!

By default, PyTorch's GELU uses an exact calculation of the function -- the hand-written code from the book uses an approximation using tanh. Luckily, you can get that same approximation from PyTorch:

         super().__init__()
         self.layers = nn.Sequential(
             nn.Linear(cfg["emb_dim"], cfg["emb_dim"] * 4),
-            nn.GELU(),
+            nn.GELU(approximate="tanh"),
             nn.Linear(cfg["emb_dim"] * 4, cfg["emb_dim"])
         )

So, training with that for 20 minutes:

giles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/
Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 4115.09it/s]
Starting rank 0 training at global step 0
  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=23,875]


Checkpoint

Continuing training
  1%|▉                                                                                                     | 307/33165 [20:00<35:39:42,  3.91s/it, loss=6.332, tps=25,142]

25,142 tokens per second -- basically the same as the non-approximate version.

So: switching to the built-in GELU made my PyTorch code run 20% faster, at about 25,000 tps rather than 21,000.

My JAX code, which used JAX's built-in GELU, ran at around 24,000 tps. I'd actually found that rather surprising, because in JAX I was training in full-fat 32-bit floating point, while in PyTorch I was using Automatic Mixed Precision (AMP) -- a special mode that allows it to use 16-bit calculations where it won't hurt the model much.

I'd found that AMP gave PyTorch a huge speedup -- from 15,402 tps to 19,797 on one test. So JAX without AMP being so much faster than PyTorch with AMP was a bit of a surprise. Its JIT is pretty amazing, but I didn't expect it to be that much faster.

Now I think that we have at least part of an explanation. I was using JAX's built-in GELU (interestingly, with its default parameters, which means that it used the tanh approximation), but the PyTorch code was using the hand-rolled one, and that unduly penalised it and erased some of the gains it got from AMP.

If I really wanted to dig into this, I suppose I might try JAX with a hand-rolled GELU to see what happened. My guess is that because of its JIT, it might actually handle it better -- the whole hand-rolled thing could be compiled into one thing on the GPU. Perhaps it would also be interesting to try the non-AMP PyTorch code with the built-in GELU. But I doubt that would really be the best use of my time (and my electricity bill), so I'll leave it here.

On the other hand, I do intend to have a look at torch.compile in the future, to see what kind of speedup I can get from it. And it might be able to compile and fuse together the hand-rolled GELU -- so that would be an interesting thing to experiment with in that post: does the built-in GELU advantage disappear if we're compiling?

But anyway, for now, lesson learned: use built-in PyTorch modules when you can. It's a pretty obvious one ;-)

[Update] On X, Sebastian Raschka noted that he used the approximate version of GELU in his code so that the models were compatible with the OpenAI weights -- they were trained with that version, so they may behave slightly differently if you use the "pure" version. That's a great point, and so I've updated my own copy of the code to use approximate="tanh".

Adding diagrams to my static site generator with D2

Tue, 25 Aug 2026 23:00:00 +0000

A lot of the time when I've been writing posts for this blog, I've felt that a diagram would really help. But they're a pain to produce well, and I think I underuse them as a result. I wanted to fix that, and wound up adding D2 support to my static site generator. I think it works pretty well!

In the past, I've tried drawing my own diagrams in LibreOffice and exporting as SVG, but my complete lack of artistic skill doesn't help:

Asking an AI to do it for me helped in simple cases:

...but with something less standard (there must be a million neural network diagrams in their training sets) it can be really fiddly to get something right.

I did some investigations into the various diagram-generating tools out there, and decided to give D2 a go. It has a simple language for specifying what your diagram should show, and the output is pretty nice:

Here's the source for that diagram:

style.fill: transparent

tokens: Tokens
tokens -> llm.token-embeddings

llm: "" {
  token-embeddings: Token embeddings
  position-embeddings: Position embeddings

  plus: "+" {
    shape: circle
    width: 36
    height: 36
    style.font-size: 24
    style.fill: transparent
  }

  token-embeddings -> plus
  position-embeddings -> plus

  input-embeddings: Input embeddings
  plus -> input-embeddings

  transformers-layers: "" {
    style.stroke-dash: 3
    style.fill: transparent
    transformers-1: Transformers layer 1
    transformers-2: Transformers layer 2
    dots: "⋮" {shape: text; style.font-size: 28}
    transformers-n: Transformers layer n
    transformers-1 -> transformers-2 -> dots -> transformers-n
  }
  input-embeddings -> transformers-layers.transformers-1

  final-norm: LayerNorm
  transformers-layers.transformers-n -> final-norm

  output-head: Output head
  final-norm -> output-head
}

output-logits: Logits
llm.output-head -> output-logits

That looks pretty clear to me!

So now, in the source for my blog posts, I have a diagrams directory. That contains subdirectories -- by convention, I create one for each post that needs diagrams -- and D2 files. These can be generated automatically when I publish:

    compile_d2_diagrams(input_path=Path(INPUT_DIR) / "diagrams", output_path=Path(NEW_OUTPUT_DIR) / "diagrams")

...

def compile_d2_diagrams(input_path, output_path):
    if input_path.is_file():
        if not input_path.name.endswith(".d2"):
            raise Exception(f"Unknown file type in D2 tree: {input_path}")
        output_path = output_path.with_suffix(".svg")
        print(f"Compiling D2 diagram in {input_path} to {output_path}")
        subprocess.check_call([
            "d2",
            "--pad=0",
            "--layout=elk",
            "--elk-nodeNodeBetweenLayers=30",
            "--elk-padding=[top=20,left=20,bottom=20,right=20]",
            input_path,
            output_path
        ])
        return
    if input_path.is_dir():
        print(f"Making diagram directory to match {input_path}: {output_path}")
        output_path.mkdir()
        for child_path in input_path.iterdir():
            compile_d2_diagrams(child_path, output_path / child_path.name)
        return
    raise Exception(f"Don't know what {input_path} is!")

(Hat tip to Evan Hahn for the with_suffix method on Path, which I wasn't aware of.)

The flags on the command line took a little bit of fiddling; the --pad=0 just gets rid of the large margins that D2 puts around the diagram by default, but the others are to tell it to use the ELK layout package with particular formatting. Its default layout has curvy lines, and I prefer the closer-to-right-angle ones that ELK provides.

Another awkward bit was in scaling; the file that is generated by that d2 command comes out pretty large (you can see it full-size here). By default, I allow images inlined into my posts to be as wide as the text, but that would still be too large here.

I use markdown2 to convert the markdown source for my posts into HTML, and there isn't any way to tell it what size an image should be using markdown-ish syntax. So for now, instead of embedding images the normal markdown way, like this:

![A simple neural network](/post-assets/neural-networks-maths/network.svg "A simple neural network")

...for these D2-generated ones I'll just embed a normal <img> tag like this:

<img src="/diagrams/adding-d2/llm-top-level.svg" alt="A GPT-2-style LLM at the top level" title="A GPT-2-style LLM at the top level" style="width: 50%">

...so that I can control the size. Perhaps more work needed there.

At some point I may go back and update my old diagrams -- at least, the really ugly hand-drawn ones -- to use this.

And a random thought: perhaps it might also make sense to include the D2 source somehow on the blog? I can imagine that it could help with accessibility in some situations, and perhaps also for any LLMs stopping by. Will have to ponder that a bit more.

What do you think? Does the D2 diagram look good to you? Or is there a better diagramming package that might work better?

Why do OpenAI's GPT-2 weights beat mine? Part four: digging into dropout

Thu, 27 Aug 2026 19:00:00 +0000

I'm still digging into a mystery about the models I've been training; although an increasing number of them beat the OpenAI GPT-2 small weights on the narrow technical measure of the loss they get on a test set, they're not as good at an instruction-fine-tuning test.

While reading about MoE models, I came across this paragraph in the Switch Transformers paper:

Our paper considers the common NLP approach of pre-training on a large corpus followed by fine-tuning on smaller downstream tasks such as summarization or question answering. One issue that naturally arises is overfitting since many fine-tuning tasks have very few examples. During fine-tuning of standard Transformers, Raffel et al. (2019) use dropout (Srivastava et al., 2014) at each layer to prevent overfitting.

So far, when running my IFT test, I'd been aiming to use the same dropout setting for the fine tune as the model concerned had used in its original pre-training. That was just because it seemed natural.

But the goal of dropout is to prevent overfitting when training over multiple epochs -- or, at least, that's how most of what I've read explains why we don't need it on modern single-epoch training runs over large datasets.

If that's the case, though, when we do multiple epochs for a fine-tune with a more restricted dataset -- exactly what I was doing for the IFT test -- it might make sense to use dropout, regardless of whether or not the model was pre-trained with it. The fine-tuning setup already tries to avoid overfitting by bailing out when a validation loss starts rising, but dropout might still help it avoid overfitting prematurely.

On the other hand, something felt a little wrong about fine-tuning a model with dropout if its pre-training had happened without it. A model pre-trained with dropout have been trained on billions of tokens, and so the model will have spent a lot of effort learning to overcome the issues that dropout causes, but one trained without it won't have that benefit. Suddenly exposing it to dropout in a much shorter fine-tuning run felt rather like asking someone who rarely drinks alcohol to take a few shots of whisky; I felt that the models might not be prepared for the effects.

As I looked into this more, I noticed another surprising thing -- there was an error in the configuration that I was using when fine-tuning the OpenAI models, both small and medium. They were originally trained with dropout (or so it's believed -- the paper doesn't say, but "Build a Large Language Model (from Scratch)" says that they were, and this config on the Hugging Face GPT-2 code agrees).

But that actually made my original puzzle of why they outperformed my models on the IFT test seem even more perplexing, at least in the light of this idea. If dropout was a good thing for fine-tuning, then so far they had been penalised by not using it -- that is, they were even further ahead of my own models than I thought they were.

It was time to take a careful look.

The fine-tunes

I fixed the config for the OpenAI weights so that my setup had dropout set to 0.1 for them, then carefully revisited the config for all of my own models, and made sure that those ones matched reality (which they did).

Now, the IFT test that I've been running has two phases:

Firstly, for each model, I run ift_generate_test_responses.py. This script trains the specified model on an IFT dataset until validation loss starts rising. It then uses the model from before that loss started going up to generate responses to a test set, and saves those responses to disk. I made a small change to it so that the dropout used in the fine-tuning phase was a required command-line parameter, with three options: model -- that is, what the model was pre-trained with -- on, which forced it to 0.1, or off, to force it to 0.
Next, I pass all of the saved test responses for all models into a second script, ift_judge.py, which sends them to an LLM judge so that each model can get a score. The script averages all scores across all questions for each model. Check the link for more details of how that script works and tries to achieve consistency across models and responses.

Now, the nice thing about the judge script was that it didn't really care whether the result files it got came from different models or the same one; it just printed out a mapping of result files to scores. So I realised I could use it to do a comparison of all models with all possible dropout settings.

For all of the models, I ran the ift_generate_test_responses.py three times, once with each of the dropout settings: model, on, and off. Then I sent all of the resulting result files -- all models, and all dropout options for each -- to the LLM judge in one go, to see what it came up with.

The results

Here are the results, consolidated into one table. For each model, I have:

Its loss on my test set -- the technical measure of quality I mentioned near the start. They're sorted by that column.
Whether or not the base training run -- the pre-train -- had dropout.
The number of fine-tuning epochs before validation loss started rising when the IFT run used a dropout setting identical to the pre-training (model).
The score that the model thus trained got from the LLM judge.
The fine-tuning epochs with dropout forced to be off.
The score for the dropout-off model.
The fine-tuning epochs for dropout forced to be on.
And finally the score for the resulting model from that.

	Test loss	Base dropout	`model` epochs	`model` score	`off` epochs	`off` score	`on` epochs	`on` score
OpenAI weights: medium	3.231442	Yes	2	42.40	2	43.75	2	42.40
JAX, overtrained one long epoch	3.324953	No	3	19.77	3	19.77	19	7.17
JAX, overtrained two normal epochs	3.326482	No	4	19.72	4	19.72	16	12.92
JAX, with MHA bias, no dropout	3.418784	No	4	18.69	4	18.69	13	13.20
JAX, no MHA bias, no dropout	3.420089	No	5	21.46	5	21.46	20	5.25
JAX, no MHA bias, with dropout	3.476802	Yes	7	17.74	5	13.22	7	17.74
OpenAI weights: small	3.499677	Yes	4	23.49	2	26.00	4	23.49
`1xrtx3090-stacked-interventions`	3.538161	No	4	13.77	4	13.77	13	14.06
`8xa100m40-stacked-interventions-1`	3.577761	No	4	10.76	4	10.76	19	7.36
Cloud FineWeb, 8x A100 40 GiB	3.673623	Yes	6	19.71	3	17.72	6	19.71
`1xrtx3090-baseline`	3.683835	Yes	6	13.53	4	15.74	6	13.53
`8xa100m40-baseline`	3.691526	Yes	4	15.15	3	14.19	4	15.15
Cloud FineWeb, 8x H100 80 GiB	3.724507	Yes	5	14.02	4	14.33	5	14.02
Cloud FineWeb, 8x A100 80 GiB	3.729900	Yes	4	11.25	3	11.34	4	11.25
Cloud FineWeb, 8x B200 160 GiB	3.771478	Yes	4	12.02	4	14.67	4	12.02
Local FineWeb train	3.943522	Yes	7	8.85	5	12.31	7	8.85
Local FineWeb-Edu extended train	4.134991	Yes	7	17.56	5	15.04	7	17.56
Local FineWeb-Edu train	4.166892	Yes	7	17.43	5	14.99	7	17.43

It's quite an intimidating wall of numbers, but there's a bunch of interesting stuff there.

Firstly: I've put the IFT score for each model where it was trained with the opposite of its pre-training dropout in bold. Let's look at the non-bold numbers first, though.

Sanity checks

If you scan down through the models, you'll see that the non-bold IFT scores -- that is, the one where the IFT test was done with the model dropout, and then the one where it was done with dropout set explicitly to the same value as the model one -- are identical in every case. That is a really reassuring sanity check. Remember, each of those numbers came from a different run of the ift_generate_test_responses.py script -- but because there is a fixed random seed, they should have been identical. They were presented to ift_judge.py in the same way as a separate model's response. The fact that it came up with identical scores tells us that it judged them as being equal, which is solid evidence for its consistency in judging results in this run (which is something that can be hard to guarantee with an LLM).

Similarly, if you look through the numbers of training epochs, the model epochs for each one matches the epochs with the dropout forced to match the model's pre-training setting, which is also reassuring -- it's certainly what you'd expect given a fixed random seed.

Epoch counts

Looking just at the on and off epochs columns, you can see something else interesting. With dropout forced to be on, the number of fine-tuning epochs is always higher than the number of epochs with no dropout, except in the case of the OpenAI medium weights and "Cloud FineWeb, 8x B200 160 GiB", where it's the same. That makes intuitive sense, I think. If you're discarding 10% of your activations when training a model, you'd expect it to take longer to converge.

But now let's look at the size of those changes. If you compare the increase in the number of epochs needed to train with dropout forced to be on, you can see that the change is much larger for those models that were pre-trained without dropout. The first of them, for example, "JAX, overtrained one long epoch", went up from 3 epochs to 19! That's way larger than, say, the change from 5 to 7 for "JAX, no MHA bias, with dropout".

That was the first indication that something interesting was happening when using dropout to fine-tune models that had been pre-trained without it.

One question is whether so many epochs on a small dataset might just be a bad idea, regardless of whether the early-stopping from validation loss helps avoid overfitting. However, way back I did some investigations into the effect of the number of epochs of training, and found that while varying it changed the results somewhat -- as you'd expect -- the effect was surprisingly small, and didn't change anything about the fundamental mystery of why the GPT-2 weights were so much better than mine. So I think we can put that aside for now.

The scores

Now let's dig into those scores. We can divide them into two groups; models that were helped by adding dropout, and models that were harmed.

In the "helped" group, we have these:

"JAX, no MHA bias, with dropout", which was pre-trained with dropout and gained 4.52 points when the IFT run used dropout.
1xrtx3090-stacked-interventions, which was pre-trained without dropout and gained 0.29 points.
"Cloud FineWeb, 8x A100 40 GiB", which was pre-trained with dropout and gained 1.99 points.
8xa100m40-baseline, which was pre-trained with dropout and gained 0.96 points.
"Local FineWeb-Edu extended train", which was pre-trained with dropout and gained 2.52 points.
"Local FineWeb-Edu train", which was pre-trained with dropout and gained 2.44 points.

In the "harmed" group, we have:

"OpenAI weights: medium", which was pre-trained with dropout and lost 1.35 points.
"JAX, overtrained one long epoch", which was pre-trained without dropout and lost 12.6 points.
"JAX, overtrained two normal epochs", which was pre-trained without dropout and lost 6.8 points.
"JAX, with MHA bias, no dropout", which was pre-trained without dropout and lost 5.49 points.
"JAX, no MHA bias, no dropout", which was pre-trained without dropout and lost 16.21 points.
"OpenAI weights: small", which was pre-trained with dropout and lost 2.51 points.
8xa100m40-stacked-interventions-1, which was pre-trained without dropout and lost 3.4 points.
1xrtx3090-baseline, which was pre-trained with dropout and lost 2.21 points.
"Cloud FineWeb, 8x H100 80 GiB", which was pre-trained with dropout and lost 0.31 points.
"Cloud FineWeb, 8x A100 80 GiB", which was pre-trained with dropout and lost 0.09 points.
"Cloud FineWeb, 8x B200 160 GiB", which was pre-trained with dropout and lost 2.65 points.
"Local FineWeb train", which was pre-trained with dropout and lost 3.46 points.

There are some patterns there, and I think that putting them into a table sorted by the score increase/decrease is a good way to visualise them:

	Base dropout	Score change
JAX, no MHA bias, with dropout	Yes	4.52
Local FineWeb-Edu extended train	Yes	2.52
Local FineWeb-Edu train	Yes	2.44
Cloud FineWeb, 8x A100 40 GiB	Yes	1.99
`8xa100m40-baseline`	Yes	0.96
`1xrtx3090-stacked-interventions`	No	0.29
Cloud FineWeb, 8x A100 80 GiB	Yes	-0.09
Cloud FineWeb, 8x H100 80 GiB	Yes	-0.31
OpenAI weights: medium	Yes	-1.35
`1xrtx3090-baseline`	Yes	-2.21
OpenAI weights: small	Yes	-2.51
Cloud FineWeb, 8x B200 160 GiB	Yes	-2.65
`8xa100m40-stacked-interventions-1`	No	-3.4
Local FineWeb train	Yes	-3.46
JAX, with MHA bias, no dropout	No	-5.49
JAX, overtrained two normal epochs	No	-6.8
JAX, overtrained one long epoch	No	-12.6
JAX, no MHA bias, no dropout	No	-16.21

One thing is pretty clear: with two exceptions, the models that were pre-trained with dropout are at the top, and the models that were pre-trained without are at the bottom.

Of the exceptions, 8xa100m40-stacked-interventions-1 is so close to "Local FineWeb train" that perhaps its position could be due to some kind of noise. 1xrtx3090-stacked-interventions is much more puzzling, however. It's a real outlier in terms of the models that were pre-trained with no dropout, with its improvement of 0.29 compared to the next closest, with a decrease of 3.4.

But if we disregard that outlier for the time being, the pattern actually does fit rather well into my original suspicion about the risks of switching on dropout when fine-tuning a model that was pre-trained without it. They really don't handle it very well!

On the other hand, it rather does put the kibosh on the idea that I based on the quote near the start of this post -- that fine-tuning with dropout is a good way to help the model learn with less risk of overfitting. In my particular case -- these specific models, this particular fine-tuning task, with this IFT data -- dropout seems to generally have a negative effect on the fine-tuning results. Even of those that were pre-trained with dropout, more than half got worse results when fine-tuned with it.

Another interesting thing that stands out from the table above is that the JAX models are at the top and the bottom. The model that was pre-trained with dropout was the one that gained the most from fine-tuning with it (or, contrariwise, lost out the most if fine-tuned without it). The models that were pre-trained without were the ones that were most harmed by being fine-tuned with.

If you look further up, at the original table of results, you'll see that the JAX models all did better than my other ones (which were trained using PyTorch) in terms of loss on my test set (the second column). I've been chalking that up to two things: the JAX models would have started their pre-training with different random initial weights, and they were all trained in full-fat float32 (unlike the PyTorch models, which used AMP). Given that I found that AMP had a negligible impact on training loss, I've been thinking that the "initial weights" aspect was the more important -- by chance, they happened to start in a place on the loss landscape with a route to a better minimum during training.

I don't think there's anything in these results that pushes against that theory, but it does suggest that there's some kind of "fragility" in the minima they have found; changing dropout from what they were pre-trained with seems to knock them out of their exceptional positions.

And finally, of course, the mystery around 1xrtx3090-stacked-interventions's anomalous position remains. I honestly don't have any theories at all about that one right now. Interestingly, it was trained with an identical configuration to our other (but less extreme) exception, 8xa100m40-stacked-interventions-1. The difference is that the first was trained on my local RTX 3090, using gradient accumulation to get a global batch size of 96, while the second was trained on a cloud machine with 8x A100 GPUs with 40 GiB each, which (using DDP) got a global batch across all GPUs of 96 without gradient accumulation. There's something going on there, but I'm not sure what.

Anyway, for now, I think it's time to wrap this one up.

Conclusion

The idea I started this post with -- that using dropout for the fine-tuning part of all of these IFT tests might be a good idea to avoid issues from the multi-epoch nature of the fine-tuning -- doesn't seem to hold up. Dropout in the fine-tuning turned out to be more often harmful than helpful, regardless of whether a model was originally pre-trained with dropout or not.

However, exactly how harmful it was seemed to be pretty strongly correlated with whether the model was originally pre-trained with dropout, the oddity of 1xrtx3090-stacked-interventions aside.

I think that while working further on solving this mystery, I should stick to not using dropout. Because adding it on for the OpenAI models made their performance worse, I think that's principled -- it's quite the opposite of making a choice to try to sweep the mystery I'm trying to solve under the carpet :-)

So that means that my task in future posts in this series is to explain this table (to go back to the format I've been using for the previous posts) -- the dropout off numbers from the table above, with rank added:

	Test loss	IFT epochs	IFT score	IFT rank
OpenAI weights: medium	3.231442	2	43.75	1
JAX, overtrained one long epoch	3.324953	3	19.77	4
JAX, overtrained two normal epochs	3.326482	4	19.72	5
JAX, with MHA bias, no dropout	3.418784	4	18.69	6
JAX, no MHA bias, no dropout	3.420089	5	21.46	3
JAX, no MHA bias, with dropout	3.476802	5	13.22	15
OpenAI weights: small	3.499677	2	26.00	2
`1xrtx3090-stacked-interventions`	3.538161	4	13.77	14
`8xa100m40-stacked-interventions-1`	3.577761	4	10.76	18
Cloud FineWeb, 8x A100 40 GiB	3.673623	3	17.72	7
`1xrtx3090-baseline`	3.683835	4	15.74	8
`8xa100m40-baseline`	3.691526	3	14.19	13
Cloud FineWeb, 8x H100 80 GiB	3.724507	4	14.33	12
Cloud FineWeb, 8x A100 80 GiB	3.729900	3	11.34	17
Cloud FineWeb, 8x B200 160 GiB	3.771478	4	14.67	11
Local FineWeb train	3.943522	5	12.31	16
Local FineWeb-Edu extended train	4.134991	5	15.04	9
Local FineWeb-Edu train	4.166892	5	14.99	10

The OpenAI small model still has a 4.54-point lead over the best of my own models, "JAX, no MHA bias, no dropout".

Previously I'd considered data quality as a possibility, and felt it was an unlikely cause. I now think I may have been premature in that, and it's worth looking into. Those two "Local FineWeb-Edu" models near the bottom were trained with sub-optimal hyperparameters and -- while they don't do super-well in this test -- they do much better than their raw test loss numbers might suggest.

But while thinking about dropout, it occurred to me that there were other levers that I'd pulled in my interventions into my original base model that might be worth investigating ¹:

Weight tying -- I honestly can't think of a reason why it might make a model better for this kind of task, but it certainly is true that the OpenAI weights use it -- while none of the ones of mine that I've been testing do. That feels worth a quick look, especially given that I have a copy of a model that I trained using it lying around.
AMP. Apart from "JAX, no MHA bias, with dropout", all of the JAX models -- trained without AMP -- did pretty well in this test (though not close to the OpenAI models). And again I have a PyTorch model that was trained without AMP on my disk somewhere, so I may as well throw it in and see how it does.
The learning rate. All of these fine-tunes are happening with a fixed learning rate of 0.00005. While I really don't want to do some kind of sweep across multiple values for all of these models, perhaps there's some way I can try to relate the fine-tuning learning rate to what the models are "used to" from pre-training and see if that helps?

So, plenty of further possibilities for this investigation. Stay tuned!

Other interventions that I decided not to check, at least at this point:
- QKV bias: all of my PyTorch models in the table apart from the two stacked-interventions ones use it, so that's been thoroughly tested.
- Weight decay: again, we have a mixture of values for that in the table and there's no obvious pattern.
- Gradient clipping: likewise.
↩

Putting my JAX-trained models on the Hugging Face Hub

Thu, 03 Sep 2026 18:00:00 +0000

I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because Transformers has been PyTorch-only since version 5 (though they say they're working to add interoperability with JAX in the future), so it would have been tough to get them working natively with AutoModelForCausalLM and the like.

But then it dawned on me that I'd already written a conversion script that could take my JAX safetensors files and convert them into ones compatible with my PyTorch code. It's actually those converted models that I use for my evals -- so I could use my existing PyTorch script to upload them.

So, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models:

"Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)"

gpjt/jax-no-mha-bias-no-dropout -- the first full LLM trained in the post, in the "Adding LayerNorm" section.
gpjt/jax-no-mha-bias-with-dropout -- the second full LLM trained in the post, in the "Dropout" section.
gpjt/jax-with-mha-bias-no-dropout -- the third full LLM trained in the post, in the "Adding bias to the MHA output projections" section.

"Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining"

gpjt/jax-with-mha-bias-no-dropout-extended -- the single-epoch, double-Chinchilla-tokens model.
gpjt/jax-with-mha-bias-no-dropout-2-epoch -- the model trained on two epochs over the Chinchilla-optimal number of tokens.

"A quick(ish) Chinchilla check"

gpjt/jax-with-mha-bias-larger-chinchilla-1 -- the slightly-larger model.
gpjt/jax-with-mha-bias-larger-chinchilla-2 -- the slightly-smaller model.

I've also added links to the posts in question.

Giles' blog

Benchmarking Qwen 3.6 35B MoE (3B active) on an RTX 3090

The model

Quantisation

Getting Llama.cpp up and running

Trying out Qwen3.6-35B-A3B -- the first quant

Trying out Qwen3.6-35B-A3B -- the second quant

The -ncmoe flag

The auto-sweep

The results

CUDA

Conclusion

Appendix: the Lem test

Why do OpenAI's GPT-2 weights beat mine?

More on the mystery

The loss landscape

OpenAI weights: better than we thought!

Probably not dropout...?

Data quality?

Overtraining?

Why do OpenAI's GPT-2 weights beat mine? Part two: the bugfix

The bug

A new baseline

Training epochs

The ranking

Conclusion

Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining

Overtraining

The experiments

The extended-train model

The two-epoch model

At last, the IFT test!

Conclusion

How I use AI on this blog

Ideation and running experiments

Writing things up, and the editorial board

AI coding

Yes to agents, but not on this blog -- yet

A quick(ish) Chinchilla check

The Chinchilla heuristic

Scaling the model

Training FLOPs

Training

Results

Conclusion

Use the built-in GELU, don't roll your own!

Adding diagrams to my static site generator with D2

Why do OpenAI's GPT-2 weights beat mine? Part four: digging into dropout

The fine-tunes

The results

Sanity checks

Epoch counts

The scores

Conclusion

Putting my JAX-trained models on the Hugging Face Hub

Trying out `Qwen3.6-35B-A3B` -- the first quant

Trying out `Qwen3.6-35B-A3B` -- the second quant

The `-ncmoe` flag