Transformer Block：RMSNorm、SwiGLU 与词表投影

从 token ids 到 logits 的结构骨架

Decoder-only 模型把 token_ids:[B,T] 通过 embedding table $E\in\mathbb{R}^{V\times H}$ 查成：

常见 pre-norm block 可抽象为：

最后一个 norm 后，vocabulary projection 产生 logits：

这是一种常见结构，不是所有 checkpoint 的唯一规范。模型可能使用 LayerNorm、post-norm、parallel residual、bias、不同 activation、shared experts 或 MoE。加载配置与权重时必须匹配具体架构，不能从“Transformer”一词推断细节。

RMSNorm：归一化 RMS，不减均值

对每个 token row $x\in\mathbb{R}^{H}$：

$g:[H]$ 是可学习 weight，沿 [B,T] 广播。与 LayerNorm 的关键边界是 RMSNorm 不计算/减去均值；若 $g=1$，输出的未加权 RMS 约为 1，但均值不必为 0。实现常用 FP32 做平方和/rsqrt 再 cast 回 activation dtype，以减小低精度归约误差；具体 accumulator 是 backend 事实。

RMSNorm 降低的不是 attention FLOPs 或 KV 容量。它提供尺度稳定性；kernel 优化机会来自减少读写、融合 residual 或 vectorized reduction。

SwiGLU：两个上投影、逐元素门控、一个下投影

设中间维度为 $I$：

它不是普通 ReLU FFN，也不是仅对同一个 projection 乘 sigmoid。Gate/up 两条独立线性路径通常可在 checkpoint/kernel 中拼成一个 [H,2I] projection，再 split、SiLU、multiply；这是 layout/fusion 实现，不改变公式。

MoE通常把这段 dense SwiGLU 替换为多个 expert SwiGLU，再由 router 稀疏选择；attention 与 residual 主干仍存在。MoE 总参数更大不意味着每 token 运行全部 experts。

Worked shape：$B=2,T=3,H=8,I=12$

X                 [2,3,8]
W_gate / W_up     [8,12]
gate / up         [2,3,12]
elementwise GLU   [2,3,12]
W_down            [12,8]
MLP output        [2,3,8]
residual output   [2,3,8]

主矩阵乘 FLOPs 约为 $2BTHI+2BTHI+2BTIH=6BTHI$，未计 activation/multiply。Prefill 的 $BT$ 较大形成 GEMM；单步 decode 的 $BT$ 小，反复读取三组权重更容易受带宽/launch 影响，详见 GEMM/GEMV 与算术强度。

Embedding 与 Vocabulary Projection

Embedding 是按 token id gather rows，不是对 one-hot [B,T,V] 做实际 dense matmul，尽管数学上等价。输入 id 必须在 [0,V) 且与 tokenizer/checkpoint vocabulary 一致；新增 special tokens 后若没有同步 resize/加载 rows，会越界或产生未训练 embedding。

输出 projection 可以独立使用 $W_{vocab}:[H,V]$，也可以 tie weights：

Tying 是参数共享语义，不代表 embedding gather 与输出 GEMM 是同一算子。部署时对 embedding/vocab 做 张量并行 切分，会让 logits 先保持 vocabulary-sharded；greedy/top-k 可用分布式归约，若下游需要全 [B,T,V] 才 All-Gather。不能无条件物化完整 logits。

可运行的组件不变量

实验问题：手写 RMSNorm 能否与 PyTorch 对齐？未加权输出 RMS 是否约为 1？SwiGLU/residual 是否保持 hidden shape？Tied vocabulary projection 是否等于 F.linear(..., embedding)？

import platform,torch
import torch.nn.functional as F

torch.manual_seed(191)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float32

def rms_norm_reference(x,weight,eps=1e-6):
    rms=torch.sqrt(x.square().mean(dim=-1,keepdim=True)+eps)
    return x/rms*weight

def swiglu(x,gate_weight,up_weight,down_weight):
    gate=x@gate_weight; up=x@up_weight
    return (F.silu(gate)*up)@down_weight

def main():
    b,t,h,i,v=2,3,8,12,16
    token_ids=torch.tensor([[1,4,7],[2,4,9]],device=DEVICE)
    embedding=torch.randn(v,h,device=DEVICE,dtype=DTYPE)/h**.5
    norm_weight=torch.randn(h,device=DEVICE,dtype=DTYPE)
    wg=torch.randn(h,i,device=DEVICE,dtype=DTYPE)/h**.5
    wu=torch.randn(h,i,device=DEVICE,dtype=DTYPE)/h**.5
    wd=torch.randn(i,h,device=DEVICE,dtype=DTYPE)/i**.5

    hidden=embedding[token_ids]
    reference=rms_norm_reference(hidden,norm_weight)
    native=F.rms_norm(hidden,(h,),norm_weight,eps=1e-6)
    torch.testing.assert_close(native,reference,rtol=1e-5,atol=1e-6)
    mlp=swiglu(native,wg,wu,wd); output=hidden+mlp
    logits_tied=output@embedding.transpose(0,1)
    logits_linear=F.linear(output,embedding)
    torch.testing.assert_close(logits_tied,logits_linear)
    assert hidden.shape==native.shape==mlp.shape==output.shape==(b,t,h)
    assert logits_tied.shape==(b,t,v)
    assert torch.equal(hidden[0,1],hidden[1,1])

    rms_before=hidden.square().mean(-1).sqrt()
    rms_after=(native/norm_weight).square().mean(-1).sqrt()
    torch.testing.assert_close(rms_after,torch.ones_like(rms_after),rtol=1e-4,atol=1e-4)
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} token_ids={tuple(token_ids.shape)} embedding={tuple(embedding.shape)}")
    print(f"hidden={tuple(hidden.shape)} gate_up={(b,t,i)} output={tuple(output.shape)} logits={tuple(logits_tied.shape)}")
    print(f"rms_norm_max_abs_diff={(native-reference).abs().max():.3e}")
    print(f"unweighted_rms_before_range=[{rms_before.min():.6f},{rms_before.max():.6f}]")
    print(f"unweighted_rms_after_max_error={(rms_after-1).abs().max():.3e}")
    print(f"tied_projection_max_abs_diff={(logits_tied-logits_linear).abs().max():.3e}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/transformer_block_reference.py。2026-08-13 在 PyTorch 2.13.0+cu130、RTX 3080 上，RMSNorm 最大差 4.768e-07，去除 weight 后 RMS 最大偏差 9.120e-06，tied projection 最大差 0。重复 token id 4 的两个 embedding rows 精确相同，验证 lookup 语义。

脚本只实现 norm + MLP + residual + vocab projection，没有复刻完整 attention block、position ids、dropout 或训练 backward。FP32 小张量也不能证明 BF16 fused RMSNorm/SwiGLU kernel 的数值或性能。

成本、失败场景与生产映射

一层 dense SwiGLU 权重 payload 主项为：

若 gate/up 合并存储，元素数不变，只改变 contiguous layout/launch。Activation 峰值取决于是否物化两条 [B,T,I]、是否 fusion、prefill chunk 和 autograd 是否关闭；推理应使用 inference/no-grad 路径，但不能只凭它推断 kernel fusion。

常见错误：

• 把 RMSNorm 写成减均值的 LayerNorm，加载 weight shape 却不报错；
• checkpoint 的 gate/up 顺序与 runtime split 顺序相反；
• 用 GELU/ReLU 替代 SiLU，或漏掉逐元素乘；
• tied weights 在加载/量化/TP conversion 后不再共享或 scale 语义不一致；
• vocabulary size、special token ids 与 embedding/logit rows 不匹配；
• 量化时对 gate/up/down 使用不受 backend 支持的 group layout，运行时 fallback；
• 认为 fusion 减少所有 FLOPs；它主要减少 launches/intermediate IO，矩阵乘主体仍存在。

模型加载必须按参数名、shape 与具体 architecture config 还原这些权重；推理量化改变它们的表示；Backend Dispatch决定实际 norm、GEMM 与 fused epilogue kernels。

参考资料

• Zhang and Sennrich, Root Mean Square Layer Normalization
• Shazeer, GLU Variants Improve Transformer
• Press and Wolf, Using the Output Embedding to Improve Language Models
• PyTorch — rms_norm
• PyTorch — embedding