gilesthomas

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 <!-- MathML: <math display="inline" xmlns="http://www.w3.org/1998/Math/MathML"><mrow><msqrt><mrow><mn>2</mn></mrow></msqrt></mrow></math> --> 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. 1 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. 2 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. 3 Promising! 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 :-) 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... --- 1. 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. ↩ 2. 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. ↩ 3. There is a table of results later on in this post where you'll be able to compare models easily. ↩

For a while I've been planning to put together a separate machine for local LLM training. Until now, I've been using my desktop PC, perry. I have an RTX 3090 installed, and can get useful training runs done (most recently, a 163M-parameter GPT-2 small style LLM in JAX), but there are a couple of problems. - perry is my daily driver. If he's doing a training run, then everything is just a little bit sluggish as CPU and GPU alike are busy. - Although I don't play games often, it's annoying to have the option ruled out for days at a time. - While the GPU is busy with a training run, I can't do other experiments in parallel -- for example, to scope out what the next step might be. And relatedly to all of those: the two-day limit to the training runs I've been doing is something I set because that's the maximum amount of time I'm willing to have perry tied up. It would be really interesting to try longer training runs! I also have longer-term plans; a multi-GPU box would be interesting to put together -- not just to have more power locally, but so that I could test larger-scale cloud multi-GPU training runs before starting to pay for expensive machines. US$15.92 an hour to rent a machine isn't a lot of money, but it adds up, especially if you're spending it while debugging parallelism issues. And finally, I've always been interested in putting together a custom water-cooling loop in a PC. I've been building my own machines since 1995 or so, but never got round to that side of things. It sounds fun! But despite all of those future plans, this is a fairly normal machine-building post -- how I repurposed an old PC, plugged in a second-hand RTX 3090 from eBay, tested it all, accidentally trained an LLM for 11 days, and almost cooked a CPU. Over time, I expect to be posting more -- and more interesting -- build details. Let's think of this as establishing the baseline. ## poppy Back before I moved to Lisbon, we had a holiday home here. When we came over, I'd bring my laptop, but that was always somewhat unsatisfactory -- limited CPU power for work, limited GPU for my occasional gaming. During Covid, we started staying in the holiday home for longer periods -- and this became too big of an annoyance to ignore. So in 2020 I put together a small form-factor PC, which I named poppy. The constraints were: - Small enough to fit in a carry-on bag. I was building the machine in London, and wanted to be able to bring her to Portugal easily, and to be able to bring her back if I wanted to. - Portable enough to quickly move around the flat. In the holiday home, the dining room was my study, so I wanted to be able to keep poppy there normally, but move her when we had guests for dinner. - Powerful enough to be able to run the games I was playing -- at the time I was a big fan of Assassin's Creed Odyssey, which didn't need a flagship card, but wasn't lightweight either. The build was a bit fiddly, like all SFF PCs. You can see the component list and build notes here on PCPartPicker, but in short she had: - An AMD Ryzen 5 3600 3.6GHz 6-Core CPU - A Noctua NH-L9a-AM4 CPU cooler - A Gigabyte X570 I AORUS PRO WIFI Mini ITX Motherboard - 32 GiB Corsair Vengeance DDR4 RAM - 2x Samsung 970 Evo 500 GB NVMe SSDs - A Zotac GTX 1660 Super 6 GiB GPU - A Lian Li PC-TU100 Mini ITX case - A Corsair SF450 450W SFF PSU She looked like this: (Gosh, I'd forgotten how... vivid our wallpaper was in that dining room.) For scale -- that case is slightly taller than two cans of coke stacked on top of each other. So, pretty small. When we moved to Lisbon full-time, I brought perry with me from London, and while he's been upgraded several times since (including adding an RTX 3090 in late 2023), he's been my daily driver since. So poppy sat in the corner of my study, sad and unused :-( It was time to bring her out again. Initial plan: get her up and running in a new, larger case, with a PSU that could potentially handle three graphics cards. ## Getting her up and running in a new case Initially, I found that she wouldn't switch on: a quick check suggested that the problem was the PSU. I'd had problems with SFF PSUs in the past, and given that the plan was to give her a new one, I just got one, along with a new, larger case -- specifically: - An ASRock Phantom Gaming PG-1600G 1600W, which would have power in spades -- an RTX 3090 goes up to about 370W at full draw, so that should hopefully handle three of them plus a CPU without problems even if one or two of the GPUs had power spikes. - A Fractal Design North XL. perry was already in a North (not the XL variant) and I love the case; the XL one looked like a good option if I was going to be cramming more GPUs in there, and had plenty of space for water-cooling. A few days later, the parts arrived. Here's a family photo: perry is to the left, poppy centre, sitting on top of her new case, and Cornélia (wearing her Flower of Shame) is to the right. For scale, Cornélia is quite a large cat. (I appreciate that that is not immensely helpful.) Time to put the old motherboard and the new PSU into the new case. Here's what it looked like: The Mini-ITX motherboard in a case designed for full ATX looks comically like a postage stamp. I switched her on, and luckily enough, everything worked! Must have been a PSU issue. The OS that she had was a more than three-year-old version of Arch, so I wiped the drives and installed the most recent version with my normal config, and it was time for a quick test. ## An LLM training run, because of course it is One of the nice things about having done all of this LLM training stuff recently is that you have a ready-made burn-in test for new hardware :-) I didn't have my JAX training code yet, but I did have the PyTorch one. Now, with her GTX 1660 Super GPU, poppy was clearly not going to be able to train an LLM of the size I could with perry's RTX 3090. I did some fiddling around with the model and training run parameters, and found that I could fit in a cut-down version of GPT-2 small with this setup: - Vocab size: 50257 -- this was fixed because I was using the GPT-2 tokeniser. - Context length: down from 1024 to 512 - Embedding dimensions: down from 768 to 512 - Number of heads: down from 12 to 8 - Number of layers: down from 12 to 8 - QKV bias: no (different to GPT-2, but the same as my own best local model). I trained it with a microbatch size of 4, gradient accumulation over 16 steps, and all other hyperparameters the same as my normal training runs on perry. The number of training tokens went down -- the model had 76,933,120 parameters, so I needed to train for just over 20x that -- about 1.5B instead of the 3.2B I've been training my other models on. I kicked that off, and out of interest, I kicked off another training run on perry with the same setup to see what happened. The perry training run went normally -- GPU running at full blast, 368W, and it completed in about 9 hours. That's less than 1/4 of the time my normal training runs take, which makes sense because time taken for this kind of thing scales roughly linearly with both the size of the model and the number of tokens, and both of those were about half the normal size. poppy was a bit more interesting. In nvtop, the GPU usage showed up as 100%, but with an "effective" utilisation of 53%. The power draw matched the latter, being 67W out of a total possible 125W. I'm not quite sure what was causing that -- clearly there was a bottleneck somewhere. Not really worth digging into, though, given that I was going to replace the card shortly. Anyway, that took 963,257 seconds to run. That's 267.57 hours, or just over 11 days. Yikes. What's kind of interesting there is that this training run not only took much longer (which is only to be expected), but that it used more electricity. 67W over 267.57 hours is just short of 18kWh, whereas 368W over 9 hours is about 3.3kWh. Buy an RTX 3090, save the planet! I decided to run my normal evals to confirm that what had come out the other end was sane. When asked to complete "Every effort moves you", perry's model said: Every effort moves you and your friends. The key is keeping you informed. But, you still need to be cautious about And poppy's said: Every effort moves you forward. Every week, you make adjustments to a certain level—just as soon as you leave your Those were actually rather good, I thought! And looking at my normal loss test confirmed that the models really weren't that bad; perry's got 3.855702, and poppy's 3.855981. That was actually better than the 3.943522 I'd got on perry before I went down my rabbit hole of optimising hyperparameters. So, that was an interesting test -- I was talking to ChatGPT about it at the time and it called it "maybe an art project", which I thought was amusing if a bit arch. Time to do something a bit more useful. ## An RTX 3090 Finding an RTX 3090 for a decent price from a trustworthy-seeming vendor is kind of hard right now. But it's still the sweet spot for price-performance if you're looking to train models locally, so I set up an alert on eBay, and eventually one popped up in Bulgaria. I bought it, and a few days later, this turned up: It's actually not as ugly as it looks in that photo -- it's considerably uglier. The stuff that looks a bit like crinkled aluminium foil is really white plastic with a kind of crystalline texture. Made me glad that I'd gone for the mesh-sided case rather than the glass one. Well, I hadn't bought it for the looks. I removed the old GTX 1660, and put in the new card, switched it on, and: /post-assets/poppy-the-training-box-1-the-beginnings/disco-3090.mp4 Wow, a disco in my PC. Lovely. ## Testing, phase 1 It was time to kick off another training run to see if it worked. This time, I did my normal GPT-2 small sized train with optimised hyperparameters. It ran for about ten minutes, and then poppy switched herself off. That didn't look good. I spent some time digging around trying to work out why my new graphics card was broken, and then happened to be sending the video above to a friend, and spotted something. Check out the Noctua fan -- the beige and brown one you can see behind the cooler mount, above the graphics card. It wasn't spinning. That's the CPU cooler fan and should always be spinning, even if slowly, when the machine is on. I log basic metrics for all of my PCs to a central InfluxDB instance, so I checked that out and: A CPU temperature spike up to about 115°C! Not good. Clearly an emergency thermal shutdown from the CPU. I initially thought that I must have knocked the fan cable loose while plugging in the new GPU -- plausible, though they were quite far apart -- but unplugging then reseating it, then powering up the machine still didn't start the fan spinning. And it was not visible in the BIOS. I then zoomed out a bit in Grafana; I only keep 30 days' worth of metrics, and it had been more than a month since I did my original burn-in test, so I didn't have anything for that. But I did have this: poppy had been idle for all of that time, and was averaging CPU temps of over 70°C. The dropoff prior to running the test was because she'd had a chance to cool down while I installed the GPU. Having spent ages setting up my InfluxDB monitoring stuff so that I have metrics for everything, I should probably actually look at them every now and then, because the fan had obviously not been doing anything for a month or so. Well, thank goodness for Amazon next-day delivery. I bought a new Noctua NF-A9x14 PWM (praying that the problem was the fan and not the header on the motherboard), and when it arrived, I put it in. ## Testing, phase 2 This time, when I powered her on, the fan was spinning. Phew. I left her running for an hour, and the CPU temperature stabilised at 35.5°C. Next, I kicked off a version of my standard LLM training run with the number of tokens reduced so that it would run for an hour. During that, the CPU temperature went up to a moderately-toasty 76°C -- not ideal, but remember that with the broken fan, she was running that hot at idle. It seemed a bit odd that it was that hot at 10% CPU usage, but given that one core was running at 100%, it didn't seem totally off. The heatsink and fan are designed for SFF PCs anyway, and those tend to run somewhat hot. The GPU temperature also went up to 70°C and stabilised there, while power draw was stably about 368W out of 370W, and GPU utilisation at 100%. That was particularly pleasing because Nvidia cards throttle at 83°C or so by default, so if I was getting a lower temperature at full power, the fans clearly had some headroom for cooling. Once that was completed, it was time for another full training run for a burn-in. ## One of these days I'll do something that doesn't involve training an LLM, I promise I kicked off my normal run. CPU and GPU temperatures stabilised at the same level as they had with the one-hour test, which was promising, so it was just a question of waiting... ...until I got this: Training complete in 144,531.113 seconds Tokens seen: 3,260,252,160 Throughput: 22,557 tokens/second Final train loss: 3.530 About 40 hours, which is pretty much standard -- certainly the same as I'd expect from perry. The smoke test: Every effort moves you and your customers to the best possible extent. Our expertise and expertise have led to our ability to Don't you just love it when your LLM tries to sell you something? 1 But anyway, loss on the test set was 3.548880, which is essentially the same as the same training run on perry too. Success! ## All done for now So, now poppy is a properly-configured training machine -- one RTX 3090, a CPU that runs a bit hot but at least doesn't do emergency shutdowns, and a case and a PSU with enough space for more GPUs. I think that the next step will be to move on to water cooling. In order to support more than one GPU, I'll need a new motherboard and probably a new CPU, so I don't think there's any point in watercooling the latter, despite its toastiness -- I'd just be buying a waterblock for it that I'd throw away in the not-too-distant future. Instead, I'll get the block for the GPU, and set up a loop to cool just that. Who knows -- maybe I can get rid of that horrendous RGB stuff at the same time! We live in hope. --- 1. Also, that "expertise and expertise" tiny model smell. ↩

I'm porting my PyTorch LLM code to JAX , using Flax as the neural network layer. For various reasons I wanted to use Safetensors to store checkpoints of the model. It took a little while to get it working; here's the trick I learned. If you look at the Safetensors docs, you'll see that it doesn't mention a JAX implementation -- indeed, searching for "safetensors jax" at the time I'm writing this gives you a link to this GitHub repo by Alvaro Bartolome -- which was last updated in 2023. However, if you look more closely at the docs, they do have a link to the Flax API . I feel this is somewhat misnamed, as it is actually a JAX API. There's no reference (again, as of the time of writing) to Flax in the source -- it's all just JAX code. And in fact Bartolome's library uses it under the hood. There is one problem, though. The API works with simple single-level dictionaries, with strings mapping directly to JAX arrays. For example, the save_file function has this signature: def save_file ( tensors : Dict [ str , Array ], filename : Union [ str , os . PathLike ], metadata : Optional [ Dict [ str , str ]] = None , ) -> None This can cause problems if you're not careful. If you look at the Flax documentation on checkpointing , it suggests that you use Orbax 1 , which has its own API and file format, but then goes on to say: When interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the nnx.State.to_pure_dict and nnx.State.replace_by_pure_dict API to convert an nnx.State to and from pure nested dictionaries. I initially put two and two together -- that and the dictionary-based API for Safetensors -- and got five, and tried feeding one of those "pure" dicts into Safetensors. I got a very confusing error: SafetensorError: dtype object is not covered It's worth digging in to why that happens. The problem is that although Safetensors is expecting a dict of strings mapping to tensors, it doesn't check that that is what it actually gets. And while the dictionaries from nnx.State.to_pure_dict are "pure", they are also nested (as the docs say!). Even for the simple model I was working with, I got a structure like this: { 'output_head' : { 'kernel' : Array ([ ... ], dtype = float32 ) }, 'token_embedding' : { 'embedding' : Array ([ ... ], dtype = float32 ) } } So, we had strings mapping to dicts, and those dicts mapped from strings to the JAX arrays. More complex models would have had deeper dict structures. Now, internally inside Safetensors, the Flax/JAX API is a simple wrapper. It iterates over the keys in the dictionary it's been provided with, and tries to convert their respective values into NumPy arrays. It does that by passing them into NumPy's asarray function, which accepts things like lists, tuples, and NumPy arrays, and converts them into arrays. JAX's own Array class exposes an interface that it recognises, so they're converted without trouble. Once it's done that, it passes the result to a lower-level Rust implementation that actually converts everything to Safetensors format. But because Safetensors didn't check types, in my case it was iterating over the top level of the dict, trying to convert the values to NumPy arrays, and got something like this: { 'output_head' : numpy . array ({ 'kernel' : Array ([ ... ], dtype = float32 )}, dtype = object ), 'token_embedding' : numpy . array ({ 'embedding' : Array ([ ... ], dtype = float32 )}, dtype = object ) } That is -- because it assumed that the values in the top-level dict were JAX Array s, it blindly tried to convert them to NumPy arrays. But they were dicts (that happened to map from strings to arrays) -- and if you ask asarray to create an array based on a random object, it happily does so and wraps that object in a NumPy array, with a dtype of object . When that is then fed into the lower-level Rust code that is trying to write the file, it encounters NumPy arrays that have a dtype it can't handle, object -- hence that error: SafetensorError: dtype object is not covered It all makes sense when you read through the code, but I was a bit perplexed for a while! I think all this might be the reason why Bartolome created his GitHub repo. In the README, he says that: There are no plans from HuggingFace to extend safetensors to support anything more than tensors e.g. FrozenDicts , see their response at huggingface/safetensors/discussions/138 . So the motivation to create safejax is to easily provide a way to serialize FrozenDicts using safetensors as the tensor storage format However, you don't need to use that library to serialise simple Flax models. Consider how PyTorch models get serialised to Safetensors; my LLMs have keys with names like out_head.weight , pos_emb.weight , and trf_blocks.0.att.out_proj.weight . They're "flat" dictionaries mapping strings to PyTorch Tensors, similar to what Safetensors wants for these Flax ones, but they use dots to separate different levels, with integers for list items and strings for field names. Looking at the pure-dict structure I had for my model: { 'output_head' : { 'kernel' : Array ([ ... ], dtype = float32 ) }, 'token_embedding' : { 'embedding' : Array ([ ... ], dtype = float32 ) } } ...you can see that you could walk the dictionary structure to generate keys like output_head.kernel and token_embedding.embedding . That would be easy enough to code up. But -- as Adithya Dsilva points out on GitHub -- you can get there even faster by using nnx.to_flat_state . That returns a (non-dict) structure like this: FlatState ([ (( 'output_head' , 'kernel' ), Param ( # 786,432 (3.1 MB) value = Array ([[ 2.3581974e-02 , 3.0957451e-02 , - 3.5088759e-02 , ... , - 4.5880198e-02 , 5.3717274e-02 , - 2.6590331e-02 ], ... , [ - 9.6302675e-03 , - 3.3276502e-02 , 5.7173111e-02 , ... , - 7.9063717e-03 , 2.0532632e-02 , 5.4753982e-02 ]], dtype = float32 ) )), (( 'token_embedding' , 'embedding' ), Param ( # 786,432 (3.1 MB) value = Array ([[ 0.00273973 , - 0.01754938 , 0.04656043 , ... , - 0.04276522 , - 0.03986642 , - 0.00781331 ], ... , [ 0.01421758 , - 0.0219186 , - 0.01701825 , ... , - 0.00793659 , 0.00500103 , 0.03839901 ]], dtype = float32 ) )) ]) If you iterate over that FlatState , you get tuples where the first element is that tuple of strings, like ('output_head', 'kernel') , and the second is a Param object wrapping the JAX Array . The tuples mirror the dot-separated string format in the PyTorch-style Safetensors files. Param objects also implement an interface that asarray can understand, so you can quickly and easily convert the FlatState to a regular dict for Safetensors: from safetensors.flax import save_file ... model_state = nnx . state ( model ) flat_state = nnx . to_flat_state ( model_state ) simple_dict = {} for tuple_key , param in flat_state : key = "." . join ( str ( key ) for key in tuple_key ) simple_dict [ key ] = param save_file ( simple_dict , "model.safetensors" ) (You need to wrap key in a str because if you have a nnx.Sequential in your model, the item in the tuple will get an integer index rather than a string). You can go the other way pretty easily too; given a model, you can load the saved checkpoint into it like this (because from_flat_state accepts raw JAX Array s in place of explicit Param s): from safetensors.flax import load_file ... simple_dict = load_file ( "model.safetensors" ) dict_flat_state = {} for key , array in simple_dict . items (): elements = key . split ( "." ) list_key = [] for element in elements : try : list_key . append ( int ( element )) except ValueError : list_key . append ( element ) dict_flat_state [ tuple ( list_key )] = array new_flat_state = nnx . from_flat_state ( dict_flat_state ) nnx . update ( model , new_flat_state ) A little more work than I'd ideally like, but given that it can be tucked away in general save_checkpoint / load_checkpoint functions, not too big a deal. Hope that's of use for other people coming across this problem! I'm beginning to feel a bit swamped with all of these libraries with names ending in -ax. It reminds me of the names of the characters in Asterix's village ... ↩