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.