FlashNorm：面向Transformer的快速归一化
机器学习
2026-04-28 v5
摘要
归一化层在大型语言模型(LLM)中屡见不鲜，却成为计算瓶颈：在具有向量与矩阵执行单元的硬件上，RMS计算阻塞后续的矩阵乘法，无法实现并行执行。我们提出FlashNorm，这是一种对RMSNorm精确 reformulation 的线性层，既通过折叠归一化权重到后续线性层中消除归一化权重，又将标量RMS归一化延迟到矩阵乘法输出后执行，从而实现两种操作的并行。此外，由于RMS的尺度不变性，一个RMSNorm后接线性层再接另一个RMSNorm可完全消除第一个RMSNorm——这一在数学上等价的简化，移除了使用QKV归一化(如Gemma~4)或在MLA模型中使用潜在归一化(如DeepSeek-V2、Mistral Small 4和OpenMythos)的注意机制前RMSNorm。相同技术也适用于LayerNorm、Dynamic Tanh(DyT)、带GLU变体的前馈网络以及基于RoPE的注意力。 在NVIDIA T4 GPU上，FlashNorm在计算受限(预填充) regime下，在SmolLM2-135M规模上实现了33-35%的延迟降低，在Llama-7B规模上实现12-14%的延迟降低。我们验证了在三个模型上的零损失权重折叠。除推理速度外，FlashNorm通过减少参数张量数量也简化了模型实现。观看我们的解释视频https://youtu.be/GEuJv34_XgU?si并访问https://github.com/OpenMachine-ai/transformer-tricks获取代码。
关键词
引用
@article{arxiv.2407.09577,
title = {FlashNorm: Fast Normalization for Transformers},
author = {Nils Graef and Filip Makraduli and Andrew Wasielewski and Matthew Clapp},
journal= {arXiv preprint arXiv:2407.09577},
year = {2026}
}