前阵子一直在跟进某个开源大模型仓库的源码，翻到归一化层的时候发现又是一个 RMSNorm。说实话，很多同学对 Transformer 里的 LayerNorm 已经足够熟了，但换成 RMSNorm 之后，往往只记住了"去掉了均值"这句话，真正问起来数学原理、为什么用 rsqrt、精度转换到底怎么转，就说不清了。这篇就顺着这个思路把 RMSNorm 完整拆一遍：从数学定义到反推公式，从 rsqrt 的硬件原理到 fp16/bf16 下的数值细节，最后给一段可以直接抄走的 PyTorch 实现和融合 kernel 的写法。
1. 为什么各家模型都默认 RMSNorm：相比 LayerNorm，它砍掉了哪一部分
先把 LayerNorm 的公式摆出来，方便后面做对比。对输入向量 x 的最后一维做归一化，n 是最后一维大小：
$$ \mu = \frac{1}{n}\sum_{i=1}^n x_i,\quad \sigma^2 = \frac{1}{n}\sum_{i=1}^n (x_i-\mu)^2,\quad \hat{x}_i = \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}} $$
归一化之后还跟着一个仿射变换：y_i = γ_i * x̂_i + β_i，其中 γ 和 β 都是可学习参数，初始值通常分别是 1 和 0。
RMSNorm 的公式就简洁很多：
$$ \mathrm{RMS}(x) = \sqrt{\frac{1}{n}\sum_{i=1}^n x_i^2},\quad \hat{x}_i = \frac{x_i}{\mathrm{RMS}(x)},\quad y_i = \gamma_i \cdot \hat{x}_i $$
对比一下能很清楚地看到，RMSNorm 砍掉了两样东西：一是均值 μ，也就是"重中心化"这一步；二是偏置项 β。只留下了"重缩放"，也就是除以均方根，然后乘可学习的 γ。
1.1 LayerNorm 的"重中心化"其实可以去掉
LayerNorm 做的事情，本质上可以分成两层：把数据拉回零均值附近，再把数据的尺度拉到一个合理范围。RMSNorm 这套做法后来的论文给过一个关键观察：对 Transformer 训练的稳定性而言，"重缩放"起主要作用，"重中心化"的贡献很有限。
怎么理解这件事？残差连接已经在帮模型解决均值漂移的问题了。Transformer 每一层输出都会跟输入做残差加法，均值信息会在这个加法的过程中被反复携带和修正，并不需要额外用一个归一化算子去强制拉均值。反而，减均值这个操作在每一层都引入了一次额外的全局规约，计算上不便宜，数值上也增加了一次减法补偿的负担。
回想一下 LayerNorm 反向传播的公式，均值 μ 会出现在每一个梯度项里，计算梯度时需要跨整个最后一维求和的量不止一个，这让融合 kernel 的实现更繁琐。RMSNorm 去掉均值之后，前向少一个 reducer，反向也少一个包含 μ 的耦合项，代码实现和 kernel 融合都会清爽很多。
1.2 数学定义与"重缩放不变性"
RMSNorm 的归一化因子是均方根而不是标准差。这两者在输入均值为零的时候是相等的，但输入均值不为零时，均方根会比标准差偏大一些。也就是说，RMSNorm 相当于"假设输入均值接近零"情况下的一种标准归一化。即便实际输入均值漂移得比较厉害，RMSNorm 依然能正常工作，因为它最关心的是把向量缩放到一个稳定尺度上。
这里有一个值得点出来的性质：重缩放不变性。如果输入 x 整体乘以一个正常数 c，那么 RMS(cx) = c * RMS(x)，归一化后的结果不变。也就是说，RMSNorm 对输入幅度整体的均匀变化不敏感，不需要像 BatchNorm 那样保存全局的 running statistics，这让它在序列长度、batch 大小动态变化的场景下天然友好。
用一个表格把两个算子的差异收拢一下：
| 对比项 | LayerNorm | RMSNorm |
|---|---|---|
| 是否需要计算均值 | 需要 | 不需要 |
| 统计量个数 | 两个（sum、sumsq） | 一个（sumsq） |
| 偏置项 β | 有 | 无 |
| 反向传播耦合项 | 包含 μ 相关项 | 只有一个全局点积 |
| 数值风险 | 用 E[x²]-μ² 可能灾难性抵消 | 直接算 E[x²]，规避该风险 |
最终效果的差距并没有那么大，毕竟决定 Transformer 稳定性的还有残差结构、注意力分数缩放、初始化策略这些因素。RMSNorm 能流行，很大程度是因为它用更少的计算代价换来了与 LayerNorm 基本相当的效果，而且在大规模训练里省下来的规约开销是实打实的。
2. RMSNorm 前向与反向的手推过程：一个公式看清所有运算
2.1 前向的三步：平方平均、倒数开方、乘系数
以单个输入行 x = [x_1, x_2, ..., x_n] 为例，RMSNorm 前向可以拆成三步。
第一步，计算均方值：
$$ ms = \frac{1}{n}\sum_{i=1}^n x_i^2 $$
第二步，计算输出缩放因子。注意这里常见的写法是：
$$ s = \frac{1}{\sqrt{ms + \epsilon}} $$
第三步，输出：
$$ y_i = \gamma_i \cdot x_i \cdot s $$
第二步和第三步其实能合并成一个表达式：y_i = γ_i * x_i * rsqrt(ms + ε)。这就是 rsqrt 优化的数学雏形。rsqrt 是 reciprocal square root，即倒数开方；它直接把"先开方再除法"变成了"一次倒数开方再全部用乘法"。这个看似微小的变化，在后面的性能分析里会体现出实际收益。
γ 的作用要单独说一下。归一化本身把所有维度拉到了同一个尺度，但模型不同维度的重要程度很可能不一样。γ 给每个维度一个可学习的缩放自由度，初始化成 1 的时候相当于不改变归一化后的分布，训练中再由优化器去调整。去掉 β 不会造成容量损失，原因是偏置项的信息经过前一层 Linear 的偏置已经能表达，归一化层再放一个 β 属于重复建模。
2.2 反向传播里隐藏的"全局规约"
反向传播是 RMSNorm 实现里最容易写错的地方。设 y_i = γ_i x_i s，其中 s = (ms + ε)^(-1/2)，记 g_i = ∂L/∂y_i，现在求 ∂L/∂x_j。
先求两个中间偏导：
$$ \frac{\partial s}{\partial x_j} = -\frac{x_j}{n s^3} $$
然后：
$$ \frac{\partial y_i}{\partial x_j} = \gamma_i s \delta_{ij} - \frac{\gamma_i x_i x_j}{n s^3} $$
代入链式法则：
$$ \frac{\partial L}{\partial x_j} = \gamma_j s g_j - \frac{x_j}{n s^3}\sum_{i=1}^n \gamma_i x_i g_i $$
注意第二项里有一个对 i 的求和：Σ γ_i x_i g_i。这意味着反向传播不是逐元素独立的，它必须跨最后一维做一次全局规约。这也是为什么 RMSNorm 的反向 kernel 比前向 kernel 难写：前向只需要一次规约，反向本质上也要一次规约，只是被藏在了梯度公式里。
一个实现上的小技巧：因为 γ_i x_i = y_i / s，所以那个点积可以改写成：
$$ \sum_i \gamma_i x_i g_i = \frac{1}{s}\sum_i y_i g_i $$
这样在 fused kernel 里可以复用已经计算出的输出 y，而不用重新读 x 和 γ。反向梯度公式整理后还能继续化简成用 y、g、s 三个量表达的版本，减少从全局内存加载数据的次数，这对访存密集的归一化算子来说影响很明显。
如果只写 PyTorch 的老实版本，直接用 autograd 就行，上述推导不会暴露。但如果你想自己写 CUDA/Triton 的反向 kernel，这个全局规约项就是第一个要处理的难点。
3. rsqrt 优化：为什么把这个除法换成乘法能提升性能
3.1 朴素实现和优化实现到底差在哪
先看两段等价代码。
朴素版本：
ms = x.pow(2).mean(dim=-1, keepdim=True) rms = torch.sqrt(ms + eps) out = weight * x / rms优化版本：
ms = x.pow(2).mean(dim=-1, keepdim=True) out = weight * x * torch.rsqrt(ms + eps)从数学上看，两者完全一样：sqrt(a) 的倒数就是 rsqrt(a)。但指令级别差别很大。精确除法在硬件上不是一条指令就能搞定的，它往往需要牛顿-拉夫森迭代或类似的高成本序列。CUDA 的浮点除法实现会经过多次乘法、减法和比较才能得到一个满足 IEEE 要求的商。把除法换成乘法之后，整个归一化过程只剩下 rsqrt 一条特殊函数指令加若干乘法，编译器也能更容易地把乘 γ、乘 x 这些操作并成 FMA 链。
当然 rsqrt 本身也不是精确数学意义上的倒数开方。GPU 上的 rsqrt.approx.f32 是硬件近似指令，精度大概在千分之一到十万分之几的水平。对归一化层来说，这个精度完全够用，因为后面还有残差、激活函数、LayerNorm 之外的各类数值操作，相对误差很快就被稀释了。但如果你的场景要求严格的可复现性，就要小心：近似 rsqrt 在不同 GPU 架构上的结果可能有细微差异，需要显式切换成精确版本。
3.2 硬件 rsqrt 的计算原理与牛顿迭代修正
硬件 rsqrt 之所以快，是因为它使用了一个足够好的初始猜测值，再配合可选的一轮迭代修正。初始猜测的逻辑可以这样理解：浮点数本身是指数加尾数结构，开方运算可以拆分成对指数减半、对尾数用一个多项式近似。GPU 里的 SFU 单元专门处理这类特殊函数，吞吐量比通用除法路径高。
如果对精度不满意，还可以手动加一轮牛顿-拉夫森迭代。假设 a = ms + ε，y 是 rsqrt(a) 的近似值，求 f(y) = 1/y² - a = 0 的根，迭代公式为：
$$ y' = \frac{1}{2} y (3 - a y^2) $$
这个迭代是二次收敛的，误差每轮变成平方级别。也就是说，如果硬件初始近似有大约 1e-4 的相对误差，一轮迭代就能压到 1e-8 以下。在训练大模型时通常不需要加，因为训练本身有随机噪声，但在某些需要精确对齐数值的推理部署场景，你会看到有人会手动写这一轮修正。
牛顿迭代用的 FMA 指令同样比精确除法便宜不少。所以"rsqrt 近似 + 牛顿修正"这种组合在性能和精度上都能兼顾，这正是很多高性能算子库里的标准套路。
3.3 rsqrt 优化值得用在什么场景
单独看一个 rsqrt 指令，性能差异可能并不夸张。真正的收益要在融合场景里看：当 RMSNorm 被实现成一个单 kernel，每个 token 的 hidden 行在寄存器里只被加载一次，pow、mean、rsqrt、乘权重这些操作全部在流水线里完成，指令数越少，整个 kernel 的耗时越短。尤其是 hidden size 很大（比如 4096 甚至 8192）时，规约和逐元素操作的占比都很高，省掉指令的效果会被放大。
但也要泼一盆冷水：在 PyTorch 的朴素写法里，torch.rsqrt对比torch.sqrt加除法的差距不会特别大，因为大的开销在pow和mean这些张量操作以及多次 kernel 启动上。rsqrt 优化的真正意义是当你要写融合 kernel、要极致压榨带宽和指令周期的时候才会体现出来。如果你的模型跑得还不错，完全没必要为了这一层优化去手写 CUDA，先保证代码正确、dtype 处理合理更重要。
4. 精度转换里的三个坑：fp16 平方、fp32 累加、eps 位置
4.1 为什么要先把 x 提到 fp32 再算统计量
大模型训练里，激活值通常是 fp16 或 bf16。如果直接在半精度下算 x_i²，会立刻踩到数值问题。fp16 的最小正常数大概是 6.1e-5，尾数只有 10 位。当激活值的量级在 1e-2 到 1e-3 时，平方之后在 1e-4 到 1e-6，已经逼近甚至低于 fp16 能精确表示的范围，很多项会变成 subnormal，有效位数严重缩水。
做个简单的实验思路：x_i = 1e-3 时，x_i² = 1e-6，这个数在 fp16 下只能用很低的相对精度表示，如果 x_i 更小，甚至直接下溢成 0。bf16 的尾数只有 7 位，相对误差更大，虽然它的指数范围和 fp32 相同，不容易下溢，但累加时的舍入噪声依然明显。
所以标准做法是：先把输入提升到 fp32，在 fp32 域内完成平方、均值、rsqrt，最后乘完权重之后再转回原来的 dtype。这段逻辑在大模型源码里几乎成了模板，不同实现之间的差异大多数只在 eps 取值和中间变量命名上。
4.2 bf16 与 fp16 的差异及累加误差
把几种常见数据格式摆一起看会更直观：
| 类型 | 指数位 | 尾数位 | 最小正常数 | 主要风险 |
|---|---|---|---|---|
| fp32 | 8 | 23 | 约 1.18e-38 | 累加维度极大时尾数溢出 |
| fp16 | 5 | 10 | 约 6.1e-5 | 平方下溢、累加精度不足 |
| bf16 | 8 | 7 | 约 1.18e-38 | 尾数太短，相对舍入误差大 |
用 fp32 做统计量计算之后，fp16 和 bf16 之间的差异就不再是 RMSNorm 的主要矛盾。比较值得留意的是累加误差。mean 操作本质上是对 n 个平方项求和再除以 n，n 是最后一维大小。如果手写 kernel 时简单地从左到右用一个变量累加，那么当 n 上千时误差会累积。PyTorch 的 CUDA reduction 会分层做部分和再合并，问题不大，但如果你自己写 CUDA 或者 Triton，规约策略就要认真设计。
对极度敏感的场景可以用 Kahan 补偿求和。示意逻辑：
total = 0.0 compensation = 0.0 for value in x_squared: y = value - compensation t = total + y compensation = (t - total) - y total = t mean = total / n这个写法在 CPU 上容易验证，在 GPU kernel 里写起来就比较繁琐，实践中大多数项目并不会启用。我自己的观点是：先保证统计量在 fp32 域，再保证规约不是单线程串行，通常就足够。
4.3 eps 加在 rsqrt 之前还是之后
这是个一眼看上去很小、实际影响输出的细节。错误写法：
out = x * (torch.rsqrt(ms) + eps)正确写法：
out = x * torch.rsqrt(ms + eps)如果 eps 加在 rsqrt 的结果上，实际效果是给每个输出分量叠加了一个常数缩放项，归一化后的分布被破坏了，eps 也根本起不到防止除零的作用。正确的位置是在开方之前让分母变成 sqrt(ms + eps)，或者直接 rsqrt(ms + eps)。
eps 的取值也要结合 dtype 看。fp32 域下 1e-6 是个常用值，很多开源模型就是这么设的；如果整个归一化都在低精度下完成，那 eps 至少要比该 dtype 能分辨的最小尺度大一些，否则相当于白设。通常建议在 fp32 域里设 1e-5 到 1e-6，然后让整个计算都走 fp32 路径，这样各类数值问题都会被隔离在统计量计算外面。
还有一个 LayerNorm 的老问题被 RMSNorm 天然绕开了：方差有时候用 E[x²]-μ² 计算，当均值远大于标准差时，两个接近的大数相减会出现灾难性抵消。RMSNorm 根本不计算均值，直接算平均平方和，不存在这一步抵消，数值路径更干净。
5. 可直接复用的 RMSNorm 代码与性能实测思路
5.1 PyTorch 实现：留出 dtype 转换的正确姿势
下面这段是经过验证的 PyTorch 版本，可以直接集成进 Transformer 模型：
import torch import torch.nn as nn class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.dim = dim self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: orig_dtype = x.dtype x_float = x.float() ms = x_float.pow(2).mean(dim=-1, keepdim=True) scale = torch.rsqrt(ms + self.eps) y = self.weight.float() * (x_float * scale) return y.to(orig_dtype)几个容易踩的细节：
mean必须加keepdim=True，否则最后一维被压缩掉，后续广播会出错。weight.float()这一行不能省。如果 weight 是 fp16 而 x_float 是 fp32，某些后端会产生隐式类型转换的额外开销，甚至引发 dtype 不匹配异常。- 最后必须
.to(orig_dtype)。RMSNorm 在模型内部通常是低精度残差流的一部分，如果这里直接输出 fp32，下游 Linear 的 autocast 行为会被打乱，显存占用也会变大。
如果你想在推理脚本里以函数形式调用，可以再封装一层：
def rms_norm_fp32(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6): x_f = x.float() ms = torch.mean(x_f * x_f, dim=-1, keepdim=True) x_normed = x_f * torch.rsqrt(ms + eps) return (weight.float() * x_normed).to(x.dtype)5.2 Triton 融合 kernel：把规约算子合并成一次遍历
如果性能敏感，或者你想写出接近线上推理实现的效果，用 Triton 写一个融合前向 kernel 更合适。核心思路是每个程序处理一行 hidden state：一次加载输入和权重，做规约、取 rsqrt、乘权重，然后一次写回。
import torch import triton import triton.language as tl from triton.language.extra import libdevice @triton.jit def rms_norm_fwd_kernel( X, W, Y, stride_x, N, eps, BLOCK_N: tl.constexpr, ): row = tl.program_id(0) cols = tl.arange(0, BLOCK_N) mask = cols < N x = tl.load(X + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32) ms = tl.sum(x * x, axis=0) / N scale = libdevice.rsqrt(ms + eps) y = x * scale * w tl.store(Y + row * stride_x + cols, y.to(X.dtype.element_ty), mask=mask)这个 kernel 的关键是把四个 Tensor 操作合并成一次数据扫描：统计量的部分和就在寄存器里，规约结果直接就地使用，中间不需要写回全局内存，也不需要重新读输入。如果 Triton 版本里libdevice.rsqrt不可用，可以退而求其次写成1.0 / tl.sqrt(ms + eps)，语义等价，只是编译器不一定会生成硬件 rsqrt 指令，应该优先用 libdevice 路径。
实践中有个额外优化点：kernel 的 grid 可以设成 batch*seq，也就是把每个 token 当一行。X 的 stride 取 hidden size，这样规约的长度就是 N。如果 BLOCK_N 比 N 大，就靠 mask 处理边界，覆盖不同 hidden size 时不用重新编译。
5.3 替换进 Transformer 模块时容易踩的坑
替换掉原有 LayerNorm，最安全的做法是保持接口一致：接收一个[batch, seq, hidden]的输入，返回同样 shape 的归一化结果。PyTorch 的 RMSNorm 类天然支持任意维度，只要最后一维等于dim就行。
权重初始化必须是torch.ones(dim)，不是随机初始化。如果你沿用其他 Linear 的初始化策略，RMSNorm 的 γ 会被打到一个不合适的尺度，前期训练可能直接不稳定。
检查输入 dtype 是否和 weight 一致。有些微调脚本里会把模型参数转成 bf16，而输入张量是 fp16，这时 forward 内部要先统一到 fp32 再输出原 dtype，不要把两个不同半精度混在一个计算链里。
性能测试时不要只盯着单算子耗时。RMSNorm 属于访存密集型算子，跟 Linear 相比，它的耗时占比通常不高。融合 kernel 的真正收益是在端到端推理中减少 kernel 启动次数和中间张量全局内存读写的总量，所以要测就用完整 attention block 或者整个模型层来做对比，单独测一个归一化算子的微 benchmark 参考价值有限。
还有一个容易忽略的事务：很多开源框架会在模型初始化后把eps写进配置里，如果你的归一化实现硬编码 eps，后续换预训练权重时可能出现数值对齐问题。最好把eps作为构造参数，并在模型配置里显式保存。
最后说点个人体会。RMSNorm 这个算子看起来简单到只有几行代码，但真把它放进大模型训练和推理链路时，数学推导、硬件指令、数值精度三块缺一不可。我在实际项目里见过最典型的故障是"loss 看起来正常，但某个中间 checkpoint 的激活分布出现尖刺"，最后定位到就是因为 RMSNorm 的统计量没有提升到 fp32，半精度下平方项下溢导致 scale 异常。所以不管你是抄现成实现，还是自己手写，都要盯住三个关键点：统计量必须在 fp32 域算、eps 必须加在 rsqrt 之前、输出 dtype 必须跟输入一致。把这三点守住，RMSNorm 就是那个看起来不起眼但极其可靠的基础算子。
