Abstract: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.
FlashNorm is mathematically identical to the original computation, it introduces no approximation and requires no retraining. The same technique extends to LayerNorm, Dynamic Tanh (DyT), feed-forward networks with GLU variants, and RoPE-based attention.
On an NVIDIA T4 GPU, FlashNorm achieves 33 to 35% lower latency on the norm-then-project operation in the compute-bound (prefill) regime at SmolLM2-135M scale, and 12 to 14% at Llama-7B scale. We verify zero-loss weight folding on SmolLM2-135M, Llama-3.2-1B, and Llama-3.1-8B.
Beyond inference speed, FlashNorm simplifies model implementations by reducing parameter tensor count, analogous to the simplification achieved by PaLM's removal of bias-parameters from all linear layers.
Watch our explainer video this https URL and see this https URL for code.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2407.09577 [cs.LG] |
| (or arXiv:2407.09577v4 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2407.09577 arXiv-issued DOI via DataCite |
From: Nils Graef [view email]
[v1]
Fri, 12 Jul 2024 00:37:55 UTC (440 KB)
[v2]
Tue, 1 Apr 2025 23:19:22 UTC (449 KB)
[v3]
Sun, 1 Jun 2025 22:12:10 UTC (584 KB)
[v4]
Wed, 22 Apr 2026 03:03:18 UTC (597 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。