Putting my JAX-trained models on the Hugging Face Hub
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-- theslightly-largermodel.gpjt/jax-with-mha-bias-larger-chinchilla-2-- theslightly-smallermodel.
I've also added links to the posts in question.
Why do OpenAI's GPT-2 weights beat mine? Part four: digging into dropout
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.
Adding diagrams to my static site generator with D2
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!
Use the built-in GELU, don't roll your own!
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.
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.
How I use AI on this blog
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.
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.
Why do OpenAI's GPT-2 weights beat mine? Part two: the bugfix
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.
Why do OpenAI's GPT-2 weights beat mine?
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.
Benchmarking Qwen 3.6 35B MoE (3B active) on an RTX 3090
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.