RMSNorm为何流行?一位开发者实测:最多只比LayerNorm快2.5%
规范化(Normalisation)是大型语言模型(LLM)中不可或缺的一部分。在GPT-2代码中,我们在多个不同位置使用了LayerNorm——在每个Transformer块中,我们在将输入的上下文向量送入多头注意力模块之前进行一次规范化,然后在将结果送入前馈网络(FFN)之前再做一次:
……此外,在通过输出头将上下文向量投影到词表空间之前,我们还会再做一次。
我们这样做是为了让模型更易于训练:通过约束进入注意力、前馈和输出头模块的向量,使其具备相对稳定的统计特性,有助于防止梯度消失或梯度爆炸。
当我尝试训练一个几乎完整的模型(这是我在用JAX对GPT-2做黑盒式重实现的一部分)而不使用LayerNorm时,训练在大约100步后完全失败——损失上升到了我预期中未初始化模型才会有的水平,并且在额外的约700步训练中一直保持在那里。(在本文后面,我会展示类似运行的结果。)
所以,规范化很重要。但它不一定非得是LayerNorm——如果你看看当前流行的开源模型,很多都使用一种叫做RMSNorm的东西,我想探究一下为什么会这样。
在这篇文章中,我将介绍LayerNorm和RMSNorm之间的差异,并展示使用每种方法训练几个GPT-2小型风格模型的结果。
原理
RMSNorm是在2019年由Biao Zhang和Rico Sennrich¹发表的论文中提出的,这篇论文出人意料地易读。他们指出,当时大多数人认为LayerNorm的一个关键部分在于均值规范化——而他们明确是在把自己的新系统与之相比较。但他们不太确定均值规范化是不是真正要紧的部分,反而觉得尺度规范化才是它真正的价值所在。
让我们深入探讨这意味着什么。首先,回顾一下LayerNorm。
LayerNorm
上面图表中的规范化层接收一批上下文向量——批次中每个序列的每个标记对应一个向量。LayerNorm对它们逐一处理²。对于每个向量,它执行两个步骤:
1. 规范化上下文向量,使其均值为零、方差为一。
2. 通过可训练的量对规范化后的向量进行缩放和偏移(又称偏置)。
让我们逐步计算。先看规范化部分;假设你从一个上下文向量开始:
(2, 4, 6)
其均值是 (2+4+6)/3 = 4。如果你从每个元素中减去它,得到:
(−2, 0, 2)
……其均值为 (−2+0+2)/3 = 0。这样就解决了均值部分。我将在下面称之为“重新居中”——希望原因显而易见!
要使方差为一,我们需要将每个元素除以整个向量的标准差(SD)(有关原因的详细信息,请参阅我之前关于LayerNorm的文章)。
标准差是方差的平方根,方差的计算方法是将每个元素与均值的差值平方,将它们全部相加,然后除以元素数量³。
现在,这里有一个巧妙的技巧。由于我们在上一步中已经将均值设为零,向量中的每个数字已经恰好等于其与均值的差值。所以我们可以直接对它们逐元素平方,并取平均值,这将给出我们的方差:
((−2)² + 0² + 2²) / 3 = 8/3 = 2⅔
同样,SD是方差的平方根,所以这里是 √(8/3) ≈ 1.63。将我们已经重新居中的数字除以该值,得到一个规范化的上下文向量:
(−2/1.63…, 0, 2/1.63…) ≈ (−1.22…, 0, 1.22…)
我们可以验证结果方差为一——记住它只是元素平方的平均值:
((−1.22…)² + 0² + (1.22…)²) / 3 = 3/3 = 1
所以,这就是LayerNorm的规范化部分。我们已经将上下文向量重新居中于其均值,并使其方差为一。
现在我们已经规范化了上下文向量,需要对其进行缩放和偏移。直觉上,这感觉有点像在撤销我们精心做的把均值归零、方差归一的工作。但值得记住的是,我们做规范化是要让结果具备稳定的统计特性;至于理想的特定属性应该是什么,那是网络需要学习的。因此,第一阶段先把一切带入一个已知的状态,然后偏移/缩放再从这个状态出发,走向网络学到的实际所需的形态。
这两个操作都使用与上下文向量大小相同的可训练向量完成——也就是说,缩放向量长度为d_emb,偏移向量也是如此。它们逐元素应用。假设我们的缩放向量是:
(3, 2, 1)
规范化后,我们的原始上下文向量是:
(−1.22…, 0, 1.22…)
所以我们将其逐元素乘以缩放向量,得到:
(−1.22… × 3, 0 × 2, 1.22… × 1) = (−3.66…, 0, 1.22…)
……缩放就此完成。类似地,如果我们的偏移向量是:
(10, 11, 12)
……那么将其应用于缩放后的向量将是:
(−3.66… + 10, 0 + 11, 1.22… + 12) = (6.34…, 11, 13.22…)
……至此,我们完成了LayerNorm。
但它进行了相当多的计算;让我们列出它们:
- 计算均值:对于长度为d_emb的上下文向量,需要d_emb−1次加法,然后一次除法。
- 重新居中于均值:d_emb次减法。
- 计算方差:我们可以使用重新居中的向量作为差值,所以这部分完成了,但我们需要所有数字的平方,这需要d_emb次乘法,用d_emb−1次加法将它们相加,然后一次除法。
- 方差到标准差:一次平方根。
- 将所有内容除以标准差:d_emb次除法。
- 按缩放向量缩放:d_emb次乘法。
- 按偏移向量偏移:d_emb次加法。
因此,对于批次中的每个上下文向量,我们总共有:
- 加法/减法:(d_emb−1) + d_emb + (d_emb−1) + d_emb = 4·d_emb − 2
- 除法:1 + 1 + d_emb = 2 + d_emb
- 乘法:d_emb + d_emb = 2·d_emb
- 平方根:1
接下来看看RMSNorm,看看它如何比较。
RMSNorm
RMSNorm背后的想法是尽量少做计算。具体来说,我们跳过均值的计算——甚至不尝试把上下文向量重新居中。相反,我们只是把它们缩放到同样的长度。为此,我们算出每个向量的平均长度——具体来说,就是均方根(RMS,当然)——然后把向量里的每个数字都除以这个值。完成后,我们执行与LayerNorm相同的缩放技巧,但不进行偏移——稍后会详细介绍。
让我们再次用数字计算;我们可以从相同的上下文向量开始:
(2, 4, 6)
其均方是,如你所料,其中所有数字平方的平均值:
MS = (2² + 4² + 6²) / 3 = (4 + 16 + 36) / 3 = 56/3 = 18⅔
均方根是它的平方根,约为4.32,所以我们除以它:
(2/4.32…, 4/4.32…, 6/4.32…) ≈ (0.46…, 0.92…, 1.38…)
你可能直观地希望结果向量的RMS为一,当然确实如此;均方是:
(0.46…² + 0.92…² + 1.38…²) / 3 = 3/3 = 1
……再次,我们对其开平方根得到RMS,再次为一。
所以你可以看到,通过除以RMS,我们将所有上下文向量重新缩放,使它们的新RMS为一。但它们的均值可能各不相同——我们没有尝试对它们做任何修复。
这是我们简化的规范化;完成后,就像LayerNorm一样,我们通过可训练的缩放向量执行逐元素缩放操作——但我们不偏移它。有趣的是,在原始论文中,他们给出了带和不带偏移两种方程——方程(4)(他们首次提出时)没有偏移,但方程(5)(他们称之为一般形式)有偏移。我不清楚他们在测试中实际使用了哪一个,但值得一提的是(感谢GPT-6 Astra为我挖掘出这一点):作者之一在GitHub上发布的代码里有一个偏置参数,默认是关闭的。
无论如何,当前的潮流似乎是使用不带偏移的RMSNorm,而这在我看来从直觉上说得通——我们通过除以RMS把上下文向量的数值尺度做了归一,所以在此基础上做“修正”以获得理想的统计特性是有意义的。但我们并没有归一化均值,因此输出的数字没有一个稳定的基线,偏移也就无从调整。当然真要加也还是可以加,但并没有直观的必要去引入一个“新均值”来替代一个本来就没被移除的旧均值。这说法可能有点模糊,但多少有点道理。
所以让我们坚持流行的无偏移版本:通过RMSNorm我们节省了多少计算?
- 计算所有数字的平方:d_emb次乘法。
- 对它们求平均以获得均方:d_emb−1次加法,然后一次除法。
- 从中获取RMS:一次平方根。
- 将所有数字除以RMS:d_emb次除法。
- 按缩放向量缩放:d_emb次乘法。
所以这给我们总计:
- 加法/减法:d_emb − 1
- 除法:1 + d_emb
- 乘法:d_emb + d_emb = 2·d_emb
- 平方根:1
与LayerNorm相比的差异:
- 加法/减法:(4·d_emb − 2) − (d_emb − 1) = 3·d_emb − 1 更少
- 除法:(2 + d_emb) − (1 + d_emb) = 1 更少
- 乘法:2·d_emb − 2·d_emb = 0 差异
- 平方根:1 − 1 = 0 差异
所以,它肯定更便宜——加法/减法操作显著减少。问题是,训练时间的节省是否值得?
在论文中,他们训练了各种类型的模型(包括一个早期的Transformer模型),并得出结论:值得。他们使用RMSNorm获得了与LayerNorm训练运行相似、甚至在某些情况下更好的损失。当然,它如今在开源LLM中的广泛使用也印证了这一点。
它对我自己的模型会有效吗?如果有提升的话,具体能得到什么样的性能改进?
是时候写一些代码并训练一些模型了。
代码
RMSNorm的PyTorch代码非常简单:
class RMSNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.emb_dim = emb_dim
self.scale = nn.Parameter(torch.ones(self.emb_dim))
self.eps = RMSNORM_EPS
def forward(self, xs):
squares = xs ** 2
mean_squares = squares.mean(dim=-1, keepdim=True)
rms = torch.sqrt(self.eps + mean_squares)
return (xs / rms) * self.scale
我们传入的上下文向量批次形状为b×n×d_emb,其中b是批次大小,n是序列长度。在forward方法中,和上面的例子一样,我们先把所有元素平方,再沿最后一个d_emb维度求平均,从而得到每个上下文向量的均值;keepdim=True用来保持维度形状不变,这样得到的是b×n×1的结果,而不是b×n。
然后我们取平方根来计算RMS,但先加一个微小的ϵ,以确保它永远不会为零(否则下一步会除以零)。最后,我们把传入的数字除以这个调整后的RMS,再乘以可训练的缩放向量,就完成了。
有一点值得注意:对我来说,在除法时加ϵ会更自然一些,像这样:
rms = torch.sqrt(mean_squares)
return (xs / (rms + self.eps)) * self.scale
不过我写的时候手边有书里的LayerNorm代码,注意到它是在开平方根之前加ϵ的,像这样:
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift
……所以我决定保持一致性。当我后来查看PyTorch内置RMSNorm类的文档时,看到它和我做的一样,所以这似乎是“正确”的做法。当数字接近零时,这两种做法差异很大,所以保持一致很重要。
既然提到了PyTorch的内置RMSNorm,那就再深挖一点。
内置PyTorch类
不久前,我注意到用PyTorch内置的GELU模块,替代Sebastian Raschka的书《Build a Large Language Model (from Scratch)》里那段手写代码后,我在本地机器上训练时获得了显著的加速——从21,000个标记/秒提高到25,000个。
PyTorch也有自己的内置LayerNorm和RMSNorm模块。
虽然我觉得这本书其实可以用PyTorch的GELU,但它展示LayerNorm代码是100%有道理的——毕竟,“Build a Large Language Model (by Bolting Together Random Bits of PyTorch)”听起来没那么顺口。同理,这组实验里我也想跑一跑自己写的RMSNorm。
但既然RMSNorm的意义就在于加速训练,那么看看这种加速相比直接使用内置模块(而非手写)有多大差距,是值得一试的。毕竟内置模块能用上融合内核(fused kernel),这可能是个巨大优势。与其像上面的Python代码那样一步一个操作——每次下沉到CUDA做计算,再回到Python确定下一步——整个计算可以被编进单个CUDA例程,一口气跑完。更妙的是,这个例程还可以整体优化:频繁使用的数据留在寄存器里、内存布局针对缓存高效利用来设计,等等。而Python实现每一步结束时,所有数据都得留在主显存(VRAM)里。这无疑是个大优势。我猜测内置GELU之所以更快,很大程度就是这个原因,所以我预期用PyTorch的规范化类会有显著的提速。
这意味着我希望代码能在不同类型的规范化之间轻松切换。
规范化切换代码
正如我在本文开头提到的,原始GPT-2代码在三个地方使用了LayerNorm:
- 每个Transformer块中两次:在MHA之前和FFN之前。
- 在最后,就在输出头之前。
Transformer部分的代码看起来像这样:
class TransformersBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
...
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
...
def forward(self, inputs):
...
shortcut = x
x = self.norm1(x)
x = self.att(x)
...
shortcut = x
x = self.norm2(x)
x, this_block_moe_routing_info = self.ff(x)
...
……输出头部分像这样:
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
...
self.final_norm = LayerNorm(cfg["emb_dim"])
...
def forward(self, in_idx):
...
x, moe_routing_info = self.trf_blocks((x, None))
x = self.final_norm(x)
logits = self.out_head(x)
...
这段代码本质上与Raschka书中的相同,只是有一些(在这篇文章的实验中未激活的)调整,以支持混合专家模型。
我想要一个配置标志——放在cfg字典里——来控制所用的规范化类型。看起来很清楚:只要想替换现有LayerNorm的各个类与它的接口兼容,就可以直接换上去。所以我将创建LayerNorm的现有行替换为:
self.norm1 = make_norm(cfg)
self.norm2 = make_norm(cfg)
...
self.final_norm = make_norm(cfg)
……并像这样定义make_norm:
def make_norm(cfg):
norm_type = cfg.get("norm_type", "handrolled_layernorm")
if norm_type == "handrolled_layernorm":
return LayerNorm(cfg["emb_dim"])
elif norm_type == "pytorch_layernorm":
return nn.LayerNorm(normalized_shape=cfg["emb_dim"], eps=LAYERNORM_EPS)
elif norm_type == "handrolled_rmsnorm":
return RMSNorm(cfg["emb_dim"])
elif norm_type == "pytorch_rmsnorm":
return nn.RMSNorm(
normalized_shape=cfg["emb_dim"], eps=RMSNORM_EPS,
elementwise_affine=True
)
elif norm_type == "no_norm":
return nn.Identity()
else:
raise Exception(f"Unknown <!--CODE_BLOCK_15744--> of {norm_type!r}")
这给了我们选项:原始的“手写”Raschka LayerNorm、我自己的RMSNorm实现、它们各自的PyTorch实现,以及一个我认为值得跑一下的无规范化选项。为PyTorch实现显式设置eps,确保它们尽可能相似⁴。我给PyTorch RMSNorm设了elementwise_affine=True,纯属谨慎——它本来就是默认值。
我最初担心的一件事是权重初始化,但后来发现这根本不是问题。在许多测试里,我发现我训练的模型仅仅因为训练前随机权重的不同,最终表现就可能差不少。如果你改变组成模型的模块——比如把LayerNorm换成RMSNorm——一旦这些模块在创建时消耗的随机性数量不同,所有后续创建的模块在初始化权重时就会拿到不同的随机数,你的测试结果就被污染了。
然而我重新看了一遍代码后,意识到这里没什么好担心的。Raschka的LayerNorm和我自己的RMSNorm都不用任何随机性,因为它们的可训练参数直接设成了1和0——PyTorch实现也是这样。我其实有点失落,因为这意味着有一段我相当得意的代码派不上用场了,不过事情确实因此简单了。⁵
这些似乎都正常——我在本地用每种设置都做了简短的测试训练,确认损失在下降。该正式测试了。
训练运行
我像往常一样,从零训练以GPT-2 small为基础的模型,有两个区别:
- 不做权重绑定(weight tying)。
- QKV矩阵上不加偏置。
这些模型大约有1.63亿个参数,非常小,但希望足够大,能展现出一些有趣的结果。
对于训练运行,我将:
- 不使用dropout。
- 给学习率做调度:前5%从零线性升温到0.0014,然后余弦衰减到0.00014。
- 训练到Chinchilla最优的约每个参数20个标记。
训练模型很花钱——以我一直在这个规模上做的模型来说,在本地RTX 3090上训练大概要2天;如果用DDP在8卡A100(每卡40 GiB显存)机器上训练,则大约花40美元(约耗时2小时30分钟)。
所以我不打算训练一大堆,决定最小合理集合是:
- 一个用旧的LayerNorm。
- 一个用PyTorch LayerNorm。
- 一个用我手写的RMSNorm。
- 一个用PyTorch的RMSNorm。
另外我还想快速跑一个无规范化测试,确认训练确实会失败。
这些都在Lambda上跑的——实例可用性出了点问题,但lambda-manager又一次证明了它很管用。我顺利启动了三个实例,把各次训练分布在它们上面。
所有测试都用惯常的命令集来跑,以确保库版本等一致:
git clone https://github.com/gpjt/ddp-base-model-from-scratch.git
cd ddp-base-model-from-scratch/
./setup_lambda.sh
uv sync
uv run torchrun --nproc_per_node=8 ddp_train.py CONFIG datasets
……CONFIG设为每个我创建的训练运行配置:
- 8xa100m40-rmsnorm-1-hand-layernorm
- 8xa100m40-rmsnorm-2-pytorch-layernorm
- 8xa100m40-rmsnorm-3-hand-rmsnorm
- 8xa100m40-rmsnorm-4-pytorch-rmsnorm
- 8xa100m40-rmsnorm-5-no-norm
各配置除了norm_type以外完全相同;无规范化那个测试还有一个小的额外差别,我们(正好)马上就来看它。
结果
先看无规范化测试。根据以往经验,我怀疑它可能会很快变糟,所以我让它每10个全局步就做一次检查点并记录训练损失,而不是每617步。果然,到第530步时它开始在每个训练步记录到一致的nan损失,于是我在第640步强制退出了。
训练损失看起来像这样:
可以看到损失从第500步开始上升,之后它干脆连损失都不记录了。梯度也冲出了图表的顶部:
所以,这算是个让人安心的结果!不出所料,训练基于Transformer的LLM而不做任何规范化,就是不行。
其他训练运行则平淡得多——它们都正常完成。我下载了模型,用惯常的留出测试集跑了一遍,得到可比较的质量指标。以下是我得到的结果:
| 模型 | 训练时间(秒) | 标记/秒 | 测试损失 |
|---|---|---|---|
| gpjt/8xa100m40-rmsnorm-1-hand-layernorm | 9,358 | 348,364 | 3.595336 |
| gpjt/8xa100m40-rmsnorm-2-pytorch-layernorm | 8,324 | 391,659 | 3.571158 |
| gpjt/8xa100m40-rmsnorm-3-hand-rmsnorm | 9,134 | 356,899 | 3.601292 |
| gpjt/8xa100m40-rmsnorm-4-pytorch-rmsnorm | 8,289 | 393,294 | 3.572383 |
链接指向Hugging Face上的模型本身。
这些结果真的很有意思!先看训练速度。
训练速度
我们这里要研究的是RMSNorm与LayerNorm的性能特征,所以先做同类比较。手写RMSNorm确实比手写LayerNorm训练快一点,但无论训练时间还是标记/秒,都只快了大约2.5%。在使用内置PyTorch模块的那组对比中,差距更小——RMSNorm只提升了0.5%。
不跑大量额外训练,我没法可靠地说这类事情的噪声水平是多少,但我的直觉(基于到现在训练过60多个这类模型的经验)是:2.5%的差异很可能是真的,0.5%则很可能只是噪声。尤其值得一提的是,两个“手写”训练恰好跑在同一台实例上,而PyTorch那两个跑在不同实例上,这让2.5%这个数字更可信。
但说到房间里的大象:融合PyTorch实现与手写实现之间的差异要大得多。LayerNorm的运行时间减少了11%;RMSNorm的加速略小一点,略低于10%。这绝对不是噪声,而且完全盖过了切换规范化类型本身带来的收益!⁶
接下来看我们模型的质量。
测试损失
同样先做同类比较。手写RMSNorm的测试损失比手写LayerNorm高出 3.601292 − 3.595336 = 0.005956,是很微小的差异。
类似地,PyTorch RMSNorm的损失比PyTorch LayerNorm高出 3.572383 − 3.571158 = 0.001225。
我倾向于认为这些数字落在噪声范围内。然而我并没有数据支撑这一点——之前研究噪声时,我测过权重初始化的变化,也测过初始权重相同、但模型初始化之后各步骤随机种子不同的运行;不过那些运行用了dropout,而本实验没有。
理论上,我的训练循环应该是完全确定性的——初始权重初始化之后不再使用任何随机性。不久前,我在8x A100机器上用相同配置跑了两次训练,相隔几个月,产出的Safetensors文件逐位完全相同。
但我们这里做的是:从一组特定的初始权重出发,然后看用不同规范化设置训练会发生什么。完全可以想象,某些初始权重对LayerNorm训练“更友好”,另一些则对RMSNorm更友好。
所以,尽管两者在RMSNorm下都略差一点这个事实很有趣,但真要获得某种确定性,我得训练一大堆不同的模型才行。眼下最好还是先归因于噪声。
现在回到那头大象……同样,这里的重大差异不在LayerNorm与RMSNorm之间,而在手写实现与PyTorch内置实现之间。LayerNorm方面,PyTorch版本损失比手写版本低 3.595336 − 3.571158 = 0.024178——更好!RMSNorm方面,差异是 3.601292 − 3.572383 = 0.028909。
这些可不算很小的差异。对照我在GPT-2代码上测试过的各种干预措施:这个差异比我切到全32位(而不是用自动混合精度AMP,0.012558)获得的改进更大,比加梯度裁剪(0.013209)更大,也比给QKV权重加偏置(0.022141)更大。
要确定这是不是噪声,我得再做大量训练运行,而现在我不打算花那份时间和金钱。感觉有两种可能:
- 它只是数值差异带来的噪声——比如某种舍入现象。想象训练刚开始时的一个微小差异:某个参数在手写代码下被更新为0.11892473,而在融合代码下是0.11892474。这会让下一步的训练结果略微不同,梯度随之不同,模型从此开始逐渐分叉。当我们以32亿个标记、33,000步的训练继续推进时,可以想象两个模型会分叉得相当远。我们是在损失地貌上颠簸下行,始终努力向低处走;起点附近的一个小扰动,就可能让最终落点改变很多。
- 或者:也许存在某种系统性因素,让PyTorch实现在这种情况下就是更好?
一方面,结果的一致性很有趣。如果PyTorch模型都优于手写模型、但提升幅度各不相同,那很容易归为偶然(我觉得概率是25%,尽管你也可以说1/6)。但两个提升幅度如此接近,就更难解释了。
另一方面,系统性差异会很奇怪,至少对我是如此。我确信手写实现和PyTorch在数学上是匹配的——不光我反复核查过代码(LayerNorm那位还是领域专家写的),我还把两者分别交给Claude Opus 5.5和GPT-6 Astra,在高思考等级下检查过,确认没有差别⁷。所以,尽管我不能完全排除错误,但为了消除它们我已经做了很多努力。
目前,我也倾向于姑且把它归为噪声。我以前见过奇怪的巧合——老读者可能记得,在我那篇关于LoRE的文章里,只在输入端使用它带来的损失增量,加上只在输出端使用它的增量,几乎恰好等于两端同时使用时的增量。这个现象在两个独立随机种子上都复现了,看起来真像那么回事——直到代码中某个本不该影响它的部分被改动,它就消失了。需要更多测试。
我认为一组有趣的实验来检验这一点:不仅换不同随机种子,还尝试完全不同的权重初始化策略。假如不用PyTorch的默认初始化,而改用GPT-2论文里描述的那种,这些结果会是什么样?(更多内容留待后续文章。)
但回到本次实验:眼下这些结果能告诉我们什么?
总结
我带着两个目标开始这组实验:理解RMSNorm的工作原理,以及看看能否复现原论文的结论——损失与LayerNorm相当,同时训练时间有可观的节省。
第一部分我很满意——如果你正在读这篇文章,希望“它如何工作”这部分对你来说也同样清楚!
损失方面,我觉得也符合预期。同样,要真正确定下来,我得做更大规模、覆盖各种初始权重的实验;但RMSNorm和LayerNorm之间的差异如此之小,这本身就令人安心。
训练速度的结果就没那么亮眼了。手写代码下,2.5%的加速看起来可能是真的;但一切换到PyTorch内置类,这点优势几乎消失。
当然,最大的差异出现在PyTorch内置类和手写类之间。融合内核带来的加速并不特别令人意外,尤其在使用内置GELU也获得类似提升之后。不过,看看用torch.compile能否缩小两者差距,肯定会很有趣。
但损失的改善绝对有意思。同样,要判断这是真实还是噪声,我还需要大量训练运行。看起来最可能是噪声;不过在两种情形下效应都是对PyTorch有利,而且幅度还这么接近,这一点始终让我放不下。
除开这些,还有一个问题悬而未决:RMSNorm并没有明显更快,尤其在使用融合内核时;而它虽然不比LayerNorm差,但也没显得更好——那么,它在这么多开源权重模型中流行,原因何在?
我首先想到的可能是:随着模型变大,它在降低成本方面的效果会更强。然而细想之后,我认为恰恰相反——随着规模扩大,它的效果反而会更小。
记住——从计算量上讲——RMSNorm为每个上下文向量省下了3·d_emb次加法/减法。这一节省会随d_emb增大而平滑扩展;当然,RMSNorm的个数也会随着模型层数增加而相应增多。所以我认为,从绝对计算量上看,总节省大约会随宽度和深度各自线性增长。
但模型其他部分的成本是随宽度超线性增长的。比如在GPT-2里,FFN把向量从d_emb维映射到4·d_emb维再映射回来,计算量按O(d_emb²)增长。更现代的实现用SwiGLU,但我相信量级是一样的。
更不用说现代模型还有更长的上下文:规范化(以及RMSNorm的节省)确实随序列长度线性增长,但全注意力在那上面是O(n²)。
所以,如果我们在这些小小的1.63亿参数模型(上下文1,024个标记)上只获得比如0.5%的改进,那么在更大的模型上,百分比改进只会更小。
即便不是如此,这个量也相当小:DeepSeek V3被广泛报道训练成本为600万美元(不含研发和一堆其他开销,但在这里用这个数字是合适的)。所以即使我们忽略上面所有分析、直接拿这个数字来算,用RMSNorm也就省下3万美元——而且这基本是上限,实际预期会明显更少。
另一种可能是,这属于架构层面的原因。这些开源模型和我这里用的GPT-2风格代码有一堆差异:它们用RoPE做位置编码,用SwiGLU做前馈网络,Mistral的注意力机制也有多处不同。表面上,我看不出这些有什么明确理由会让RMSNorm更快;但我当然可以想象,它们可能“偏爱”某一种规范化、从而带来更低的损失。也许这就是人们实际遇到的情况。
但可能性还有很多。如果大家普遍发现它跟LayerNorm一样好,那也许它就只是种时尚——某个模型用了它,下一个模型也没有特别理由改回去,如此延续。此外,LLM领域还有一个普遍趋势是去掉偏置(比如QKV权重上的偏置),RMSNorm正好契合这一趋势。而且,它确实更简单一点!
无论原因是什么:如果RMSNorm只是通过跳过重新居中这一步而比LayerNorm简单了一点点,那会不会有更简单的东西、而且效果更好?在和Claude Opus 5.5讨论这个话题时,它提到了《Transformers without Normalization》⁸——一篇非常有意思的论文,其中一位作者在GitHub上有一篇很棒的导读。他们注意到,规范化层在训练中学到的变换,最终看起来有点像tanh函数那种S形映射,于是干脆尝试直接用tanh来替代,并把这套方法称为Dynamic Tanh(DyT)——结果很不错!
这篇文章的第一句话是“规范化是LLM中不可或缺的一部分”。这话错了吗?我想得看你对“规范化”怎么定义——你能把他们的tanh技巧看作某种规范化吗?——虽说Yann LeCun是那篇论文的合著者,我大概不该贸然不同意他对规范化的定义;-)
不过,他们的DyT系统更适合放到后续文章里讲。至于现在:RMSNorm在模型质量上似乎与LayerNorm基本相同,而且对我一直在玩的那类模型,它可能还略快一点——至少在不使用融合模块的情况下。
但用PyTorch融合模块替代我们手写的规范化,在速度上是巨大的胜利。本实验中随之而来的损失改善也很有意思,不过我当然得先做大量额外实验,才敢说它是真的!
目前,我想说:我很高兴学到了RMSNorm的工作原理,但我不打算把它作为训练运行的标配加进去。进一步的实验,就堆到日后要研究的几百件事里去吧。等我学会了RoPE、SwiGLU以及所有其他内容之后……
感谢阅读!
---
脚注:
1. Biao Zhang, & Rico Sennrich. (2019). Root Mean Square Layer Normalization. ↩
2. 有一种更早的规范化技术叫BatchNorm,它跨整个批次操作,但似乎已经不再受欢迎。论文中有一些关于原因的讨论。↩
3. 或者如果我们使用贝塞尔校正(Bessel's correction),则除以元素数量减一。我在这篇文章中使用不带校正的版本,但数学上无论哪种方式都成立。↩
4. Raschka的LayerNorm实现不使用贝塞尔校正,这使其与原始GPT-2代码兼容;具体来说,它用了unbiased这个参数:
var = x.var(dim=-1, keepdim=True, unbiased=False)
根据PyTorch的var文档,这是PyTorch 2.0之前表示以下内容的方式:
var = x.var(dim=-1, keepdim=True, correction=0)
根据LayerNorm文档:
> 方差通过有偏估计器计算,相当于torch.var(input, correction=0)。
所以这很好地匹配。↩
5. 以防我将来想用它,这里贴出来:
from contextlib import contextmanager
...
@contextmanager
def stable_rng():
after_seed = random.getrandbits(63)
try:
yield
finally:
random.seed(after_seed)
torch.manual_seed(after_seed)
torch.cuda.manual_seed_all(after_seed)
有了它,你可以把任何可能从各个随机数生成器中抽取可变数量随机性的代码,包进with stable_rng():块里,之后无论发生什么,状态都能保证稳定。在我早期的一些代码里,为了确保不同代码路径无论模型配置如何都抽取相同量的随机性,我费了不少周折;用这个会容易得多。↩
6. 我想到的一种可能性:LayerNorm获得更多加速,可能是因为手写代码中的低效已在PyTorch实现里得到解决。记住,我们最初用的代码是教学性质的,用来解释LayerNorm如何工作,所以它这样做:
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift
现在,回想上面解释LayerNorm时我提到的:一旦你把向量重新居中,就可以借用它来计算方差,因为重新居中的向量本身就是一组与均值的差值。所以你可以更高效地把上面改写为:
mean = x.mean(dim=-1, keepdim=True)
recentered = x - mean
var = (recentered ** 2).mean(dim=-1, keepdim=True)
norm_x = recentered / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift
这样可以省下两次均值和差值的计算,所以有望提速。
或者,当然,它也可能只是噪声,或者完全是别的东西;-) ↩
7. 我本来想在“我让LLM检查了一下”那里加一句防御性的“FWIW”,但如今,对这类事情,我觉得没这个必要。↩
8. Zhu, J., Chen, X., He, K., LeCun, Y., & Liu, Z. (2025). Transformers without Normalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). ↩