A quick(ish) Chinchilla check
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.
Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining
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.
Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)
This post is the capstone of the most long-running series on my blog. In December 2024 (!), I started reading Sebastian Raschka's book "Build a Large Language Model (from Scratch)", and worked through it carefully. Being who I am, despite trying to apply a strict "no side quests" policy, I found myself zooming off and digging into all kinds of things.
It's time to wrap it up. I had decided that the endpoint would be to build and train an LLM from scratch just using my notes -- no reference to the book, no reference to the model code I'd written when following the book. After an X/Twitter poll, I decided to use JAX for that, just to make sure that I really was building it from scratch and not regurgitating bits of PyTorch code like a bad coding LLM spitting out half-digested lumps of Stack Overflow.
In my last post, I showed how I built a JAX training script that mirrored what I had built for the original PyTorch version of the model. To test it as I went along, I used it to train a really dumb "LLM", which instead of trying to predict the next token for every token in an input sequence, instead predicted the input -- that is, if you fed it
The fat cat sat on the mat
It would return the same thing. I called that an A-to-A model.
In this post, I'll show you how I turned it into a GPT-2 model, and then trained it from scratch on my RTX 3090 (using the parameter counts for the original paper's "small" size). What turned out really well with this is that I found a route that meant that almost every component I added made the model better! That's not guaranteed -- sometimes different aspects of an AI model depend on each other, so adding A without also adding B makes things worse. But (admittedly with a bit of backtracking in places) I was able to find a route that shows a nice clear progression.
The final training run took 37 hours 15 minutes -- compared to 40 hours, 38 minutes for an equivalent PyTorch model. That is despite it being full-fat 32-bit -- the PyTorch one was using Automatic Mixed Precision (AMP), which allowed it to use 16-bit calculations in places where it would be relatively harmless in terms of loss.
When asked to continue "Every effort moves you", it came back with a decent response:
Every effort moves you closer to your goals, but if you are unsure of what it takes, you don’t
The model got 3.418784 loss on my held-back test dataset, as compared to my PyTorch model's 3.538161, and even more impressively, it was better than the original GPT-2 small's result of 3.499677 on the same dataset! However, just as I found previously, the OpenAI weights still beat mine consistently in instruction fine-tuning challenges.
Let's get started.
Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run
For over a year, I've been using Sebastian Raschka's book "Build a Large Language Model (from Scratch)" -- and the multitude of side-projects that have branched out from reading it -- as something like a curriculum for learning about modern AI. The one final task I had set myself was to build and train an LLM from scratch just using my notes -- no reference to the book, no reference to the model code I'd written following the book.
As an output, I wanted something as good as my best PyTorch model based on Raschka's code -- a base model, trained on 3.2B tokens, that my (admittedly limited) evals ranked as being close to the original GPT-2 small's quality.
I wanted to use a different framework, just to make sure I wasn't parroting code that I'd somehow memorised, so I asked people on Twitter which one I should use, and the winner was JAX.
I took a slightly different route to Raschka's book; he takes an inside-out perspective, explaining things like attention, gradually building up a complete GPT-2-style model, and then building a training loop on top of it. I wanted to go outside-in: I'd put together a training harness to train the simplest-possible model with an API similar to a real LLM, get that working to my satisfaction, and then add features to that simple model, one by one, until it had the full architecture in place. The plan (which actually worked out nicely!) was that I'd be able to show how each change improved things.
That's all done now, and I'm posting about it in two parts; in this one, I'll explain how I built the training harness, and in the next, I'll show the actual building and training of the LLM.
So let's get started!
Flax debugging: making a hash of things
I was debugging an issue with a JAX/Flax NNX training loop the other day, and found a neat little trick to help debug it. Specifically, I wanted to see if the issue was with my model, my loss function, my optimiser settings, or the "plumbing" of the training loop itself -- were gradients actually coming through and being applied to the parameters?
I could print out the loss and the gradients, but printing out the parameters to see if they were changing was unhelpful -- any given update might only change a small number of parameters, or might change them such a small amount that I'd not notice -- especially given that the model had 77 million of them!
Let's take a look.
JAX: commitment issues
Imagine you have JAX code like this, and run it on a machine with CUDA set up:
key = jax.random.key(42)
cpu0 = jax.devices("cpu")[0]
with jax.default_device(cpu0):
array = jax.random.randint(
key,
(530640, 6, 1024),
0, 50_000,
dtype=jax.numpy.uint16
)
array.block_until_ready()
item = array[0]
item.block_until_ready()
We're creating a big array, blocking until it's ready (JAX is asynchronous, so this makes sure that it's actually finished creating it), then getting the first item, and as a belt-and-braces thing making sure that that is ready too. How long do you think those last two lines -- a simple retrieval of a 6 x 1024 array from a larger one -- will take? Some tiny fraction of a second would seem reasonable.
But running it on my machine just now, the answer is a bit of a surprise: just over 5 seconds. And if you try to
get array[1] immediately afterwards, it still takes about 1.2s. Further lookups into
array consistently take more than a second -- so while the larger
initial number might be something to do with setup -- maybe internal stuff being JITted -- that's clearly not the whole story.
Something is making these seemingly-simple array lookups take much longer than you'd
expect them to.
Let's dig into that.
JAX backends and devices
There's nothing like writing your own code with a framework to clarify how things
fit together! Continuing with my port of my PyTorch LLM code to
JAX, I wanted to load up a large dataset:
the 10,248,871,837 16-bit unsigned integers in the train split of
gpjt/fineweb-gpt2-tokens.
That's just over 19GiB of data.
from safetensors.flax import load_file
...
full_dataset = load_file(dataset_dir / f"train.safetensors")["tokens"]
When I ran that, I got a CUDA out-of-memory error:
jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 19.09GiB.
That makes sense! The allocation it was trying to do is exactly the size of the data I was trying to load. I have an RTX 3090 with 24 GiB, but some is already used up by the OS, various apps, and a model that the code creates earlier on.
But in PyTorch land, I was used to things being loaded into RAM by default, and only moved over to the GPU when I asked it to do that. JAX was clearly loading to the GPU by default. How could I stop it from doing that for this case? The load into the GPU was happening inside Safetensors, in code I couldn't directly control.
Understanding how to do it helped me understand a little bit more about JAX.
Using Safetensors with Flax
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.
On first looking into JAX
Much have I travell'd in the realms of gold,
And many goodly states and kingdoms seen;
Round many western islands have I been
Which bards in fealty to Apollo hold.
Oft of one wide expanse had I been told
That deep-brow'd Homer ruled as his demesne;
Yet did I never breathe its pure serene
Till I heard Chapman speak out loud and bold:
Then felt I like some watcher of the skies
When a new planet swims into his ken;
Or like stout Cortez when with eagle eyes
He star'd at the Pacific -- and all his men
Look'd at each other with a wild surmise --
Silent, upon a peak in Darien.
John Keats, On First Looking into Chapman's Homer
I've been working with PyTorch quite a lot for the last couple of years, and feel like I've come to a reasonably solid understanding of how it all fits together. Working through Sebastian Raschka's book "Build a Large Language Model (from Scratch)", training my own LLMs locally and in the cloud, rebuilding Andrej Karpathy's 2015-vintage RNNs -- over time, it all adds up!
But, of course, there are other frameworks, and one I kept hearing about was JAX. While it's less dominant than PyTorch, it has a reputation for a certain cleanliness, a certain purity. And having spent time over the last couple of weeks working through the tutorials, and translating small PyTorch examples into it, I've been really impressed.
In this post I want to give an overview -- to report back to beginners like me, still living in PyTorch-land, on my new discovery. Less like Herschel discovering Uranus, and more like a 16th-century European coming back after having discovered something that the people who lived there were perfectly well aware of. What is this JAX thing, and how does it differ from PyTorch?