FlashNorm: 用于Transformer的快速规范化
FlashNorm: Fast Normalization for Transformers
AI总结：
FlashNorm通过消除归一化权重和并行执行RMS规范化，提升Transformer模型的计算效率，减少延迟并简化实现。
摘要：
归一化层在大型语言模型（LLMs）中普遍存在，但代表了计算瓶颈：在具有不同向量和矩阵执行单元的硬件上，RMS计算会阻塞后续的矩阵乘法，阻止并行执行。我们提出了FlashNorm，即RMSNorm后接线性层的精确改写，其（i）通过将归一化权重折叠到后续的线性层中消除归一化权重，（ii）将标量RMS归一化推迟到矩阵乘法的输出，使两个操作能够并行执行。此外，由于RMS的标量不变性，RMSNorm后接线性层再接另一个RMSNorm允许第一个RMSNorm被完全消除——这是一种数学上等价的简化，可以移除使用QKV归一化的模型中的预注意力RMSNorm（例如Gemma 4）以及具有潜在归一化的MLA模型（例如DeepSeek-V2、Mistral Small 4和OpenMythos）中的归一化。同样的技术也扩展到LayerNorm、动态双曲正切（DyT）、具有GLU变体的前馈网络以及基于RoPE的注意力。在NVIDIA T4 GPU上，FlashNorm在计算受限（预填）领域，在SmolLM2-135M规模下，规范后投影操作的延迟降低了33-35%，在Llama-7B规模下降低了12-14%。我们验证了三个模型上的零损失权重折叠。除了推理速度外，FlashNorm还通过减少参数张量数量简化了模型实现。
英文摘要：
Normalization layers are ubiquitous in large language models (LLMs) yet represent a compute bottleneck: on hardware with distinct vector and matrix execution units, the RMS calculation blocks the subsequent matrix multiplication, preventing parallel execution. We present FlashNorm, an exact reformulation of RMSNorm followed by a linear layer that (i) eliminates the normalization weights by folding them into the subsequent linear layer, and (ii) defers the scalar RMS normalization to the output of the matrix multiplication, enabling the two operations to execute in parallel. Additionally, by the scale invariance of RMS, an RMSNorm followed by a linear layer followed by another RMSNorm allows the first RMSNorm to be eliminated entirely -- a mathematically identical simplification that removes the pre-attention RMSNorm in models using QKV-normalization (e.g., Gemma~4) and in MLA-models with latent normalization (e.g., DeepSeek-V2, Mistral Small 4, and OpenMythos). The same techniques extend to LayerNorm, Dynamic Tanh (DyT), feed-forward networks with GLU variants, and RoPE-based attention. On an NVIDIA T4 GPU, FlashNorm achieves 33 - 35% lower latency on the norm-then-project operation in the compute-bound (prefill) regime at SmolLM2-135M scale, and 12 - 14% at Llama-7B scale. We verify zero-loss weight folding on three models. Beyond inference speed, FlashNorm simplifies model implementations by reducing parameter tensor count. Watch our explainer video https://youtu.be/GEuJv34_XgU?si and see https://github.com/OpenMachine-ai/transformer-tricks for code.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
