Use the built-in GELU, don't roll your own!

www.gilesthomas.com · 2026-08-20 · by Long Nguyen ·generic
Switching from the hand-rolled GELU module I'd been using to PyTorch's built-in one unsurprisingly made training models faster -- but I was surprised to see it was 20% faster.

Use the built-in GELU, don't roll your own!

Writing the post that I wished I'd found when I started learning whatever it was...

About

Contact

Archives

Categories

Blogroll

+ nn.GELU(approximate="tanh"), nn.Linear(cfg["emb_dim"] * 4, cfg["emb_dim"]) ) `

So, training with that for 20 minutes:

`giles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/ Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 4115.09it/s] Starting rank 0 training at global step 0 0%| | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=23,875]

Checkpoint

Continuing training 1%|▉ | 307/33165 [20:00<35:39:42, 3.91s/it, loss=6.332, tps=25,142] `

25,142 tokens per second -- basically the same as the non-approximate version.

So: switching to the built-in GELU made my PyTorch code run 20% faster, at about 25,000 tps rather than 21,000.

My JAX code, which used JAX's built-in GELU, ran at around 24,000 tps. I'd actually found that rather surprising, because in JAX I was training in full-fat 32-bit floating point, while in PyTorch I was using Automatic Mixed Precision (AMP) -- a special mode that allows it to use 16-bit calculations where it won't hurt the model much.

I'd found that AMP gave PyTorch a huge speedup -- from 15,402 tps to 19,797 on one test. So JAX without AMP being so much faster than PyTorch with AMP was a bit of a surprise. Its JIT is pretty amazing, but I didn't expect it to be that much faster.

Now I think that we have at least part of an explanation. I was using JAX's built-in GELU (interestingly, with its default parameters, which means that it used the `tanh`approximation), but the PyTorch code was using the hand-rolled one, and that unduly penalised it and erased some of the gains it got from AMP.

If I really wanted to dig into this, I suppose I might try JAX with a hand-rolled GELU to see what happened. My guess is that because of its JIT, it might actually handle it better -- the whole hand-rolled thing could be compiled into one thing on the GPU. Perhaps it would also be interesting to try the non-AMP PyTorch code with the built-in GELU. But I doubt that would really be the best use of my time (and my electricity bill), so I'll leave it here.

On the other hand, I do intend to have a look at `torch.compile` in the future, to see what kind of speedup I can get from it. And it might be able to compile and fuse together the hand-rolled GELU -- so that would be an interesting thing to experiment with in that post: does the built-in GELU advantage disappear if we're compiling?

But anyway, for now, lesson learned: use built-in PyTorch modules when you can. It's a pretty obvious one ;-)

\[Update\] On X, Sebastian Raschka noted that he used the approximate version of GELU in his code so that the models were compatible with the OpenAI weights -- they were trained with that version, so they may behave slightly differently if you use the "pure" version. That's a great point, and so I've updated my own copy of the code to use `approximate="tanh"`.

www.gilesthomas.com · 5 images · 195 links

Comments (0)

No comments yet