{"slug": "use-the-built-in-gelu-don-t-roll-your-own", "title": "Use the built-in GELU, don't roll your own!", "summary": "PyTorch's built-in GELU function is 20% faster than a hand-rolled version when training GPT-2 small models, according to a developer's benchmark. The same code training the same model on the same data ran at about 21,000 tokens per second with the hand-coded GELU and about 25,200 tokens per second with PyTorch's built-in version, a 20% throughput increase. The discovery suggests that the hand-rolled GELU was consuming about 17% of training time, and the developer plans to switch to the built-in version.", "body_md": "Unsurprisingly, [PyTorch's own built-in GELU function](https://docs.pytorch.org/docs/2.13/generated/torch.nn.GELU.html)\nis faster than the hand-rolled one I've been using to date. But I *was* surprised at\nhow *much* faster using it made things when training my models. I discovered this\naccidentally just now while working on something unrelated, but am logging the details\nhere for anyone else that might find it useful.\n\nThe headline numbers: the same code, training the same model on the same data, ran at about:\n\n`approximate=\"tanh\"`\n\n, which\nuses 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.\n\nAnd what is particularly\ninteresting is that there aren't that many GELUs going on -- it's a GPT-2 small-style\nmodel, with 12 layers. So that's 12 GELUs handling tensors\nshaped `(batch_size, seq_len, 4 * d_emb)`\n\n, which is `(6, 1024, 3072)`\n\nfor my training setup.\nGiven that the rest of the model is doing all of the normal full attention stuff for GPT-2, it's\n*really* surprising that the GELUs alone must have been taking up so much of the time. The throughput\nnumbers mean that we must have been spending about 17% of our time on the extra overhead from the hand-rolled\nversion, so that sets a lower bound for how much time the GELUs were taking up.\n\nMore info below the fold.\n\nBack when I was doing the \"interventions\" part of my [LLM from scratch series](/llm-from-scratch), training\ndozens of GPT-2 small-sized models in the cloud and on my local machines, to keep things simple I used\nthe original model code from Raschka's book.\nThat happens to have its own implementation of the GELU function --\nyou can see [my copy here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/16dd249447754c3a7be4f1212ccaf312c7416d92/gpt.py#L78).\n\nI'm not that sure why the hand-rolled version is in there -- he covers the maths, but the specific implementation isn't explained in that much depth, and it seems rather like boilerplate, just a \"type this in and use it\" kind of thing. By contrast, for example, while he does explain the maths behind cross-entropy loss in similar detail, we use the built-in function for it rather than coding it up ourselves.\n\nWhen I switched to using JAX for [my own from-scratch implementation](/2026/07/llm-from-scratch-34b-building-and-training-gpt-2-small-in-jax),\nI decided to not bother porting the boilerplate, and just used [JAX's own built-in version](https://docs.jax.dev/en/latest/_autosummary/jax.nn.gelu.html#jax.nn.gelu).\n\nI was revisiting the PyTorch code -- I'm in the process of extending it with mixture-of-experts support, about which more in a later post -- and decided to switch from the hand-written GELU to the PyTorch one just to tidy things up a bit. I noticed something interesting -- my new MoE code suddenly seemed to speed up.\n\nWas that a mirage? Or had I discovered part -- or even all -- of the reason why the JAX code was so much faster than the PyTorch code? With PyTorch, I was typically getting training speeds of about 21,000 tokens per second, while in JAX I was getting 24,000 tps or so. I'd been chalking that up to JAX's JIT compilation, but could it have been just a result of a random implementation choice I'd made?\n\nI did three partial test training runs, letting each one run for 20 minutes to allow the training speed to settle down from any startup overhead. Firstly, with the old hand-coded GELU:\n\n``` bash\ngiles@poppy:~/Dev/ddp-base-model-from-scratch ((HEAD detached at 16dd249))$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/\nFetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3173.90it/s]\nStarting rank 0 training at global step 0\n  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=19,982]\n\nCheckpoint\n\nContinuing training\n  1%|▊                                                                                                     | 257/33165 [20:07<42:53:57,  4.69s/it, loss=6.570, tps=20,920]\n```\n\nSo it was getting 20,920 on average over those 257 global steps. That speed was in line with [the original run](/2026/04/llm-from-scratch-32k-interventions-training-our-best-model-locally-gradient-accumulation#the-new-local-baseline)\nof the configuration I was using.\n\nNext, I introduced the built-in PyTorch GELU with no arguments:\n\n``` python\n-\n-class GELU(nn.Module):\n-\n-    def forward(self, x):\n-        return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2.0 / torch.pi)) * (x + 0.044715 * torch.pow(x, 3))))\n-\n-\n-\n class FeedForward(nn.Module):\n\n     def __init__(self, cfg):\n         super().__init__()\n         self.layers = nn.Sequential(\n             nn.Linear(cfg[\"emb_dim\"], cfg[\"emb_dim\"] * 4),\n-            GELU(),\n+            nn.GELU(),\n             nn.Linear(cfg[\"emb_dim\"] * 4, cfg[\"emb_dim\"])\n         )\n```\n\nThat does the full calculations for GELU, rather than using the `tanh`\n\n-based approximation that\nthe hand-rolled code did. After 20 minutes, it looked like this:\n\n``` bash\ngiles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/\nFetching 4 files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 11222.22it/s]\nStarting rank 0 training at global step 0\n  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=23,953]\n\nCheckpoint\n\nContinuing training\n  1%|▉                                                                                                     | 307/33165 [20:00<35:40:59,  3.91s/it, loss=6.331, tps=25,134]\n```\n\nSo this time we were getting 25,134 tokens per second -- 20% faster!\n\nBy default, PyTorch's GELU uses an exact calculation of the function -- the hand-written\ncode from the book uses an approximation using `tanh`\n\n. Luckily, you can get that same\napproximation from PyTorch:\n\n```\n         super().__init__()\n         self.layers = nn.Sequential(\n             nn.Linear(cfg[\"emb_dim\"], cfg[\"emb_dim\"] * 4),\n-            nn.GELU(),\n+            nn.GELU(approximate=\"tanh\"),\n             nn.Linear(cfg[\"emb_dim\"] * 4, cfg[\"emb_dim\"])\n         )\n```\n\nSo, training with that for 20 minutes:\n\n``` bash\ngiles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/\nFetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 4115.09it/s]\nStarting rank 0 training at global step 0\n  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=23,875]\n\nCheckpoint\n\nContinuing training\n  1%|▉                                                                                                     | 307/33165 [20:00<35:39:42,  3.91s/it, loss=6.332, tps=25,142]\n```\n\n25,142 tokens per second -- basically the same as the non-approximate version.\n\nSo: switching to the built-in GELU made my PyTorch code run 20% faster, at about 25,000 tps rather than 21,000.\n\nMy 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.\n\nI'd [found](/2025/12/llm-from-scratch-28-training-a-base-model-from-scratch) that\nAMP gave PyTorch a huge speedup -- from 15,402 tps to 19,797 on one test. So JAX without AMP\nbeing so much faster than PyTorch with AMP was a bit of a surprise. Its JIT is pretty\namazing, but I didn't expect it to be *that* much faster.\n\nNow I think that we have at least part of an explanation. I was using JAX's\nbuilt-in GELU (interestingly, with its default parameters, which means that it used the `tanh`\n\napproximation), but the PyTorch code was using the hand-rolled one, and that unduly\npenalised it and erased some of the gains it got from AMP.\n\nIf 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.\n\nOn the other hand, I do intend to have a look at `torch.compile`\n\nin the future, to see\nwhat kind of speedup I can get from it. And it might be able to compile and fuse together\nthe hand-rolled GELU -- so that would be an interesting thing to experiment with in that\npost: does the built-in GELU advantage disappear if we're compiling?\n\nBut anyway, for now, lesson learned: use built-in PyTorch modules when you can. It's a pretty obvious one ;-)", "url": "https://wpnews.pro/news/use-the-built-in-gelu-don-t-roll-your-own", "canonical_source": "https://www.gilesthomas.com/2026/08/built-in-gelu", "published_at": "2026-08-20 01:12:45.983231+00:00", "updated_at": "2026-08-20 01:12:48.046166+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools"], "entities": ["PyTorch", "GPT-2", "JAX", "Raschka"], "alternates": {"html": "https://wpnews.pro/news/use-the-built-in-gelu-don-t-roll-your-own", "markdown": "https://wpnews.pro/news/use-the-built-in-gelu-don-t-roll-your-own.md", "text": "https://wpnews.pro/news/use-the-built-in-gelu-don-t-roll-your-own.txt", "jsonld": "https://wpnews.pro/news/use-the-built-in-gelu-don-t-roll-your-own.jsonld"}}