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 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.
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. 1
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?
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
jax-gpt2-chinchilla
here.jax-gpt2-2x-chinchilla
.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 , giving us:
...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.
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!
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 4 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!
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.
Time to put this all together.
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 5.
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 :-)
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.
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. ↩