前言

nanoGPT 笔记结尾列举了一些优化新组件：RoPE、RMSNorm、SwiGLU、GQA。本文接着记录RMSNorm、SwiGLU、GQA。有了这几个组件，一个 nanoGPT 风格的 vanilla decoder block，就能升级成 LLaMA / Qwen 风格的现代 block。

---

一、先看结论：四处换件（五分钟极简版）

1.1 四组件表

从 vanilla Transformer（2017 / BERT / nanoGPT）升级到现代 LLM（LLaMA / Qwen）的四处关键换件：

组件	旧（vanilla / BERT / nanoGPT）	新（现代 LLM）	解决
位置编码	学习式绝对 PE（wpe 查找表）	RoPE	位置编成旋转角，天然相对、可外推到更长序列
归一化	LayerNorm（减均值、带 bias；2017 原版/BERT 还是 post-LN）	pre-RMSNorm	去掉均值中心化和 bias，更省同稳；norm 统一放子层之前
FFN 激活	两层 MLP（线性 + ReLU/GELU + 线性）	SwiGLU	加"门控"：一路当阀门筛另一路，同算力下表达更强
注意力	MHA（每个头独立 Q/K/V）	GQA	多个 Q 头共享一份 K/V，推理时 KV cache 省几倍显存

1.2 一段话概括

Vanilla Transformer 到现代 LLM 换了四处组件。位置编码从学习式绝对 PE 换成 RoPE，把位置编成旋转角，天然相对、可外推。归一化从 LayerNorm 换成 RMSNorm，去掉了减均值和 bias，只保留按均方根缩放，更省同稳；位置上现代模型统一用 pre-norm（2017 原版/BERT 是 post-norm，GPT-2/nanoGPT 已经是 pre-norm）。FFN 从两层 MLP（ReLU/GELU）换成 SwiGLU，多引入一路"门控"逐元素相乘，同等算力下表达力更强（为了参数量对齐，隐藏维从 4d 缩到约 8/3 d）。注意力从 MHA 换成 GQA，让若干个 query 头共享同一份 K/V 头，推理时 KV cache 显存砍到几分之一，几乎不掉点。

---

二、RMSNorm：把 LayerNorm 砍掉一半

2.1 LayerNorm

nanoGPT 每个子层前都有一个 LayerNorm。对一个 d 维向量 $x=(x_1,\dots ,x_d)$，LayerNorm 沿特征维做四步，

$$\mu =\frac{1}{d}\sum _{i=1}^d{x}_i,{\sigma }^2=\frac{1}{d}\sum _{i=1}^d(x_i-\mu {)}^2$$

$$y_i=\underset{\text{标准化}}{\underbrace{\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}}}\cdot \underset{\text{缩放}}{\underbrace{{\gamma }_i}}+\underset{\text{平移}}{\underbrace{{\beta }_i}}$$

拆成四个动作：① 算均值 $\mu$；② 减去均值（re-centering，重新居中）；③ 除以标准差（re-scaling，重新缩放）；④ 可学习的仿射 $\gamma ,\beta$。

nanoGPT 的实现：

class LayerNorm(nn.Module):
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))          # γ
        self.bias   = nn.Parameter(torch.zeros(ndim)) if bias else None  # β
    def forward(self, x):
        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)

2.2 RMSNorm：只保留"按均方根缩放"

RMSNorm（Root Mean Square Norm，Zhang & Sennrich 2019）认为 LayerNorm 里真正起作用的是"除以尺度"（re-scaling），“减均值”（re-centering）这一步贡献不大，可以整个删掉。 于是它不算均值、不减均值、也不要 bias，只用 均方根（RMS） 做分母：

$$\mathrm{RMS}(x)=\sqrt{\frac{1}{d}\sum _{i=1}^d{x}_i^2+ϵ},y_i=\frac{x_i}{\mathrm{RMS}(x)}\cdot {\gamma }_i$$

对照 LayerNorm：没有 $\mu$、没有减法、没有 $\beta$。 分母从"标准差（要先减均值再平方）“变成"均方根（直接平方，不减均值）”。

实现极简：

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))     # 只有 γ,没有 β
        self.eps = eps
    def forward(self, x):
        # x.pow(2).mean(-1) 是均方;rsqrt 直接给出 1/RMS
        inv_rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (x * inv_rms) * self.weight

省下：一次求均值、一次逐元素减法、一半的仿射参数（去掉 β）。在大模型里 norm 被调用成千上万次，这点节省累积起来相当可观，而实测下游效果 几乎无损，归一化的核心价值在于把激活的 尺度 拉回稳定区间（防止某层输出爆炸/消失，梯度好传），这件事 re-scaling 能干，re-centering 可有可无。

2.3 pre-norm

现代 block 都是 pre-norm，即 x = x + Sublayer(Norm(x))，norm 作用在进入 attention/MLP 之前，而不是 2017 原版的 post-norm（x = Norm(x + Sublayer(x))）。

pre-norm 的好处：残差通路是干净的恒等映射，梯度能从顶层直通底层，几十上百层也能稳定训练，不需要精细的 learning-rate warmup。 现代做法把两件事合起来 pre-RMSNorm（LLaMA、Qwen）。

post-norm → pre-norm 发生在 2017 原版/BERT → GPT-2 ；
LayerNorm → RMSNorm 发生在 GPT-2 → LLaMA 。

---

三、SwiGLU：给 FFN 装一道"门"

3.1 原版 FFN

nanoGPT 的 MLP（FFN）就是"升维→激活→降维"：

$$\mathrm{FFN}(x)=W_2\cdot \mathrm{GELU}(W_{1}x)$$

其中 $W_1:d\to 4d$（升维 4 倍），$W_2:4d\to d$（降回）。代码：

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd)   # W1, 升维
        self.gelu   = nn.GELU()
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)   # W2, 降维
    def forward(self, x):
        return self.c_proj(self.gelu(self.c_fc(x)))

所有信息都被同一个激活函数处理后往下走。

2017 原版 Transformer 的 FFN 用的是 ReLU，这里用的是 GELU。
下面将 ReLU 到 SwiGLU 演化线简单梳理下。

3.2 σ（sigmoid）函数

$$\sigma (z)=\frac{1}{1+e^{-z}}\in (0,1)$$

S 形曲线：$z\to -\infty$ 时趋于 0，$z\to +\infty$ 时趋于 1，$\sigma (0)=0.5$，处处光滑可导（${\sigma }^{\prime }=\sigma (1-\sigma )$）。它扮演门控开合度：把任意实数压到 0~1，当作"允许多少信息通过"的软阀门。

3.3 第一步：ReLU → GELU

起点是 ReLU：

$$\mathrm{ReLU}(x)=\max (0,x)$$

优点是正区间梯度恒为 1，计算只是一个比较，负区间置零带来稀疏性。但有两个问题：① 落入负区间的单元输出 0、梯度也是 0，一旦长期为负就学不动；② 在 0 处有硬拐角，响应是"全有或全无"的硬开关。

GELU（Hendrycks & Gimpel 2016）给出了软化方案，

$$\mathrm{GELU}(x)=x\cdot \Phi (x)$$

$\Phi (x)$ 是标准正态分布的 CDF（$\Phi (x)=P(Z\le x),Z\sim \mathcal{N}(0,1)$）。直觉来自随机正则化：每个输入 $x$ 以概率 $\Phi (x)$ 被保留、以概率 $1-\Phi (x)$ 被置零（“输入越大越可能活下来”），这个随机操作的期望就是 $x\cdot \Phi (x)$。由于 $\Phi$ 和 $\sigma$ 都是 S 曲线、形状接近，有常用近似

$$\mathrm{GELU}(x)\approx x\cdot \sigma (1.702x)$$

这样负区间不再是零，$x$ 略小于 0 时 GELU 输出一个小负值、梯度非零，dead neuron 缓解；整条曲线处处光滑，优化更平顺。2018 年 GPT-1 和 BERT 先后采用 GELU，从此它成为 NLP Transformer 的默认激活，GPT-2（nanoGPT）沿用至今。

3.4 统一视角：这些激活全是"值 × 门"

把每个激活都写成"值 × 门"的形式：

阶段	写成"值 × 门"	门的性质	门和值的关系
ReLU	$x\cdot H(x)$（$H$ 为阶跃函数）	硬门（0/1 突变）	同源（门由 $x$ 自己决定）
GELU	$x\cdot \Phi (x)$	软门（0~1 连续）	同源（自门控）
SiLU / Swish	$x\cdot \sigma (x)$	软门，与 GELU 形状几乎一样	同源
GLU	$(W_{v}x)\odot \sigma (W_{g}x)$	软门	独立（门有自己的参数 $W_g$）
SwiGLU	$(W_{up}x)\odot \mathrm{SiLU}(W_{gate}x)$	光滑非单调门	独立

3.5 第二步：GLU——把门独立成一条通路

GLU（Gated Linear Unit，Dauphin et al. 2016）给信息配一道门控。把输入通过两个独立的线性层，一路当"值（value）“，另一路过 σ 当"门（gate）”，两者 逐元素相乘，

$$\mathrm{GLU}(x)=(W_{v}x)\odot \sigma (W_{g}x)$$

注意这里增加了学习参数，之前是"每个数自己决定过不过"，现在是"网络用一套参数 $W_g$ 看完整个输入向量后，决定值路 $W_{v}x$ 的每一维放行多少"。门和值解耦，门控面对不同的 token，gate 会算出不同的开合模式。类似 LSTM 的门机制。

3.6 第三步：SwiGLU = 用 SiLU 当门的 GLU

SwiGLU（Shazeer 2020）把门的激活从 σ 换成 SiLU（$z\cdot \sigma (z)$，即 Swish${}_{\beta =1}$）。SiLU 门不像 σ 那样饱和在 (0,1)，还带一段线性成分和负区间的小谷（非单调），梯度更通畅。于是 FFN 变成 三个 权重矩阵（gate、up、down），

$$\mathrm{SwiGLU}\text{-}\mathrm{FFN}(x)=W_{\text{down}}(\underset{\text{门}}{\underbrace{\mathrm{SiLU}(W_{\text{gate}}x)}}\odot \underset{\text{值}}{\underbrace{W_{\text{up}}x}})$$

代码：

import torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden = swiglu_hidden(config.n_embd)              # ← 见 3.7,不是简单的 4d
        self.w_gate = nn.Linear(config.n_embd, hidden, bias=False)  # 门
        self.w_up   = nn.Linear(config.n_embd, hidden, bias=False)  # 值
        self.w_down = nn.Linear(hidden, config.n_embd, bias=False)  # 降维
    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

Shazeer 在论文里系统对比了 ReGLU / GEGLU / SwiGLU 等 GLU 变体替换 FFN 的效果，SwiGLU / GEGLU 实证 loss 最低，文中说"我们无法解释为什么这些结构有效，把成功归于神的眷顾。" PaLM（2022）、LLaMA（2023）先后采用，SwiGLU 成为现代 LLM 的默认 FFN。

演化时间线如下

年份	事件
2010/2012	ReLU 进入深度网络（Nair & Hinton / AlexNet）
2016	GELU 提出（Hendrycks & Gimpel）；GLU 提出（Dauphin et al.）
2017	原版 Transformer 用 ReLU；Swish 被架构搜索发现
2018	GPT-1、BERT 改用 GELU——Transformer 里 GELU 替代 ReLU 的时间点
2020	Shazeer 提出 SwiGLU/GEGLU；T5 v1.1 用 GEGLU
2022–2023	PaLM、LLaMA 用 SwiGLU，成为现代 LLM 默认 FFN

3.7 隐藏维从 4d 缩到 8/3 d

原版 FFN 只有 2 个矩阵，SwiGLU 有 3 个。如果隐藏维还用 4d，参数量和计算量会凭空多出 50%，对比不公平。为保持 参数量大致不变，LLaMA 把隐藏维缩小：

• 原版 FFN 参数：$d\times 4d+4d\times d=8d^2$
• SwiGLU 参数（3 个矩阵，隐藏维 $h$）：$dh+dh+hd=3dh$
• 令 $3dh=8d^2\Rightarrow h=\frac{8}{3}d\approx 2.67d$

就是，换 GLU 类激活，多了个矩阵，为了参数不变，隐藏层要乘 2/3。

def swiglu_hidden(d, multiple_of=256):
    h = int(2 * (4 * d) / 3)                       # 8/3 d
    return multiple_of * ((h + multiple_of - 1) // multiple_of)  # 向上取整对齐

---

四、GQA：让多个 Q 头共享一份 K/V

4.1 MHA 的推理瓶颈：KV cache

标准 MHA（Multi-Head Attention） 把 Q、K、V 各切成 $h$ 个头，每个头有 自己独立 的 K、V。

自回归生成时，每步只新增一个 token：第 t 步的新 token 要和前 t−1 个 token 的 K、V 做注意力。如果每步都把整段历史的 K/V 从头重算一遍，代价是 $O(T^2)$ 级的浪费。所以推理框架把每一步每层每个头算好的 K、V 缓存下来复用，这就是 KV cache。它的显存占用：

$$\text{KV cache}\propto 2\times \underset{\text{头数}}{\underbrace{h}}\times \underset{\text{每头维度}}{\underbrace{d_{head}}}\times \underset{\text{序列长}}{\underbrace{T}}\times \underset{\text{层数}}{\underbrace{L}}$$

长上下文（大 $T$）+ 多头（大 $h$）时，KV cache 会吃掉巨量显存。注意瓶颈在 K/V 的头数 $h$ 上，而不在 Q，Q 是当步现算现用的，不需要缓存。

4.2 MQA：省到极致（但掉点）

MQA（Multi-Query Attention）让所有 Q 头共享同一份 K/V（K/V 只有 1 个头）。KV cache 直接缩小 $h$ 倍，推理飞快。代价：表达力受限，质量有可感知的下降，训练也可能不稳。

4.3 GQA：在两个极端之间取平衡

GQA（Grouped-Query Attention，Ainslie et al. 2023）是 MHA 和 MQA 的折中：把 $h$ 个 Q 头分成 $g$ 组，每组共享一份 K/V。于是 K/V 有 $g$ 个头（$1\le g\le h$）：

MHA (g = h):   Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
               K1 K2 K3 K4 K5 K6 K7 K8     ← 8 个 Q 头,8 份 KV(最贵)

GQA (g = 2):   Q1 Q2 Q3 Q4 | Q5 Q6 Q7 Q8
               \___KV_a___/   \___KV_b___/  ← 8 个 Q 头,只 2 份 KV(省 4 倍)

MQA (g = 1):   Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
               \________ KV_0 _________/     ← 8 个 Q 头,共 1 份 KV(最省)

MHA 和 MQA 是 GQA 的两个端点：$g=h$ 退化成 MHA，$g=1$ 退化成 MQA。LLaMA-2/3 70B 用 $h=64$、$g=8$，KV cache 省 8 倍，质量几乎与 MHA 持平。

4.4 为什么省显存、又几乎不掉点

• 省：KV cache $\propto g$ 而非 $h$，选 $g=h/8$ 就省 8 倍显存 + 省对应的 K/V 投影和访存。
• 不掉点：Q 头数不变（注意力的"提问多样性"由 Q 决定，一个没少），只是让相邻几个 Q 头"看同一套 K/V"。经验表明多个 Q 头共享 K/V 的信息损失很小，配合少量继续训练几乎可以完全补回。用极小的质量代价，换极大的推理省用，是它成为现代标配的原因。

注意训练是 teacher-forcing、整个序列一次并行 forward，每个 batch 的 K/V 当场算、当场用、用完即弃，不存在"跨步复用"，所以 GQA 最大的收益在训练阶段不存在。

4.5 代码骨架

核心是repeat_kv。Q 投影出 $h$ 个头，K/V 只投影出 $g$ 个头；算注意力前，把 K/V 每头 复制 $h/g$ 份，对齐到 $h$ 个头：

def repeat_kv(x, n_rep):
    # x: (B, T, n_kv_head, head_dim) -> (B, T, n_kv_head * n_rep, head_dim)
    if n_rep == 1:
        return x
    B, T, n_kv, hd = x.shape
    return (x[:, :, :, None, :]
              .expand(B, T, n_kv, n_rep, hd)
              .reshape(B, T, n_kv * n_rep, hd))

class GroupedQueryAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_head    = config.n_head        # Q 头数 h
        self.n_kv_head = config.n_kv_head     # KV 头数 g (g | h)
        self.n_rep     = self.n_head // self.n_kv_head
        self.head_dim  = config.n_embd // config.n_head
        self.wq = nn.Linear(config.n_embd, self.n_head    * self.head_dim, bias=False)
        self.wk = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)  # 更小
        self.wv = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)  # 更小
        self.wo = nn.Linear(config.n_embd, config.n_embd, bias=False)

    def forward(self, x, freqs_cis):
        B, T, C = x.shape
        q = self.wq(x).view(B, T, self.n_head,    self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_head, self.head_dim)   # 只有 g 个头
        v = self.wv(x).view(B, T, self.n_kv_head, self.head_dim)
        q, k = apply_rotary_emb(q, k, freqs_cis)   # ★ RoPE 在这里旋转 Q、K
        k = repeat_kv(k, self.n_rep)               # 把 g 个 KV 头复制成 h 个
        v = repeat_kv(v, self.n_rep)
        # 之后照常:先 transpose(1,2) 成 (B, h, T, hd),
        # 再 scores = q@k^T/sqrt(hd) -> +causal mask -> softmax -> @v -> wo
        ...

4.6 RoPE 和 GQA

• RoPE 旋转的是 每个 K 头向量本身，旋转角只取决于该 token 的 位置 $m$，不关心这个 K 头被哪些 Q 头共享。
• GQA 共享的是 K/V 头这个物件。一个被 4 个 Q 头共享的 K 头，先按自己的位置旋转好，再被复制给 4 个 Q 头用，它们拿到的是同一个已经旋转过的 K，位置语义完全一致，不冲突。

顺序上，先 apply_rotary_emb（旋转 Q/K）→ 再 repeat_kv（复制 K/V 头）。 两个操作正交，RoPE 管"位置"，GQA 管"共享"。

---

五、常见问题

1. “RMSNorm 是效果更好，还是只是更省？”
   主要是 更省 + 同等稳，效果基本打平（有时略好）。它的价值不在"提点"，而在"用更少的计算拿到一样的稳定性"。

2. “GQA 会不会掉点？掉多少？”
   会掉一点点，但极小。经典结论：GQA（如 8 组）质量非常接近 MHA，远好于 MQA（1 组），而 KV cache 省得和 MQA 一个量级。所以它是"性价比最高的折中"。

---

总结

一条主线：一长（RoPE）一稳（RMSNorm）一强（SwiGLU）一省（GQA）。大多数换件在削冗余，SwiGLU 在加表达，判据永远是工程目标。