PyTorch 内置 GELU 替代手写版，训练快 20% · 干货攻略

• 链接: https://x.com/rasbt/status/2090254612557156587
• 分类: x-tips
• 来源: X @rasbt
• 作者: Jay
• 更新: 2026-08-23

这是什么

PyTorch 内置的 torch.nn.GELU 比手写 GELU 快约 20%，实测从 ~21,000 tok/s 提升到 ~25,000 tok/s（RTX 3090，GPT-2 small，AMP 模式）。这是 Giles Thomas（@gpjt）在 2026 年 8 月 20 日发布的博客文章的核心发现，rasbt（Sebastian Raschka）在 X 上推荐后引发关注。

为什么值得关注

谁在分享，解决了什么问题

@rasbt（Sebastian Raschka，LLMs-from-scratch 作者）在 X 上转发了 @gpjt 的博文，并预告将把这一发现写进《Build a Large Language Model (From Scratch)》教材的 bonus 章节。

问题背景： rasbt 的书从零实现 GPT 类 LLM，代码中手写了一个 GELU 激活函数（tanh 近似版本）。这个手写版本在教学中清晰展示了 GELU 的数学原理，但在生产训练场景下存在显著性能开销。

核验后的关键数据

测试环境（来自 Giles Thomas 博客原文）：RTX 3090，PyTorch AMP 模式，GPT-2 small（12 层），tensor shape (6, 1024, 3072)，训练 20 分钟后取稳定值。

实测结论：即使切换到 PyTorch 内置的 approximate="tanh" 模式（数学上与手写版完全等价），速度也与精确版几乎相同，说明 PyTorch 的 C++/CUDA 实现本身就有大幅优化，手写 Python 版本无论如何优化都追不上。

核验过程

官方来源 1：PyTorch 文档（torch.nn.GELU）

来源：https://docs.pytorch.org/docs/2.13/generated/torch.nn.GELU.html

PyTorch 2.13 确认了 nn.GELU 有两个模式：

• approximate='none'（默认）：精确计算 GELU(x) = x * Φ(x)，其中 Φ 是标准正态分布的 CDF。
• approximate='tanh'：tanh 近似版本 GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))

与 rasbt 书中的手写实现数学等价（后者用 torch.pow(x, 3)，前者用 torch.tensor(2.0 / torch.pi) 做系数）。

官方来源 2：rasbt/LLMs-from-scratch 仓库

来源：https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/gpt.py

仓库中 ch04/01_main-chapter-code/gpt.py 确认手写 GELU 实现：

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))

用于 FeedForward 层：

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),                          # ← 手写类
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

官方来源 3：Giles Thomas 博客实测数据

来源：https://www.gilesthomas.com/2026/08/built-in-gelu

实测为单卡 RTX 3090，PyTorch 的 torchrun 分布式训练脚本，使用 AMP（Automatic Mixed Precision）。三个 20 分钟实测：

1. 手写 GELU → 20,920 tok/s
2. nn.GELU() → 25,134 tok/s
3. nn.GELU(approximate="tanh") → 25,142 tok/s

交叉验证结论

• rasbt 在 X 帖子中确认：书中使用近似 GELU 是刻意选择，目的是与 OpenAI 预训练权重兼容（GPT-2 使用的是 tanh 近似版本训练）。
• 近似版与精确版在速度上几乎无差异，说明优化收益来自 PyTorch 内置实现的底层（C++/CUDA）优化，而非算法差异。
• rasbt 预告 LLMs-from-scratch 会更新 bonus 章节（本攻略核验时 X 帖子本身未提供该章节链接，故记为"原帖主张，未完全核验"）。

上手步骤

步骤 1：替换手写 GELU

在 rasbt 书中，手写 GELU 类通常定义在 gpt.py 里，只需两处改动：

改动前：

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))

改动后（推荐）：

# 删除手写 GELU 类，直接用 PyTorch 内置
# nn.GELU()  默认是精确版
# nn.GELU(approximate="tanh")  等价于原手写版数学公式

步骤 2：修改 FeedForward 层引用

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
-           GELU(),
+           nn.GELU(),                    # 精确版，推荐
+           # 或 nn.GELU(approximate="tanh")  # 等价原手写版数学
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

步骤 3：验证速度提升

# 原有训练命令不变，速度应该直接提升约 20%
uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/

观察 tps（tokens per second）数值：旧版 ~21k tok/s → 新版 ~25k tok/s。

步骤 4：精度验证（可选）

若用于从零预训练新模型，建议跑几个 epoch 对比 loss 曲线，确认精度无退化：

import torch.nn as nn

# 精确版（GELU 论文原版）
gelu_exact = nn.GELU()           # GELU(x) = x * Φ(x)

# tanh 近似版（与 rasbt 书手写版数学等价）
gelu_approx = nn.GELU(approximate="tanh")

x = torch.randn(2, 10)
assert torch.allclose(gelu_exact(x), gelu_approx(x), atol=1e-3)

坑与适用边界

⚠️ 精度差异：对已预训练权重有影响

rasbt 本人指出：GPT-2 等模型使用 tanh 近似版训练，如果你要从零预训练新模型，两种模式都可以用。但如果是加载已有预训练权重做推理/微调，应使用 approximate="tanh" 以保证激活值对齐，否则可能有轻微数值偏差。

推荐做法：

⚠️ Giles Thomas 实测是单卡 GPT-2 small

20% 提升来自 RTX 3090 + AMP + 单卡配置。多卡、混合精度关闭、或不同 GPU 架构，收益比例可能有差异。

⚠️ 手写 GELU 仍有教学价值

rasbt 书中的手写实现清晰展示了 GELU 的数学原理（tanh 近似的动机、0.044715 系数的来源），学习阶段推荐保留手写版理解原理，上线训练时切换到内置版。

⚠️ torch.compile 可能弥合差距

Giles Thomas 在博客末尾提到：后续计划测试 torch.compile()，有望把手写 GELU 编译融合进 CUDA 图，从而缩小与内置版的差距。如果你的项目已启用 torch.compile，GELU 替换的收益可能变小——这个方向值得跟进。

一句话结论

把 nn.GELU() 或 nn.GELU(approximate="tanh") 换掉手写的 GELU 类，训练速度立涨约 20%；从零训新模型用精确版，加载 GPT-2 预训练权重则用 approximate="tanh" 保持对齐。