SLaNC: Static LayerNorm Calibration

TL;DR

The paper tackles FP16 overflow during LayerNorm in Transformer inference by introducing SLaNC, an analytical method that computes LayerNorm input scales offline using only static preceding weights. This eliminates the need for calibration data or runtime scaling while guaranteeing stable, overflow-free computations and preserving FP32-like accuracy. The approach yields scaling factors for standard MLP blocks, Llama-style MLP blocks, and attention blocks, derived from Frobenius-norm expressions of weight matrices. Experimental validation on Llama models demonstrates that SLaNC prevents FP16 overflow, aligns FP16 results with FP32 baselines on perplexity metrics, and supports low-latency inference across diverse hardware. Overall, SLaNC offers a practical, theoretically justified solution for robust FP16 Transformer inference without runtime overhead.

Abstract

The ever increasing sizes of Large Language Models (LLMs) beyond hundreds of billions of parameters have generated enormous pressure on the manufacturers of dedicated hardware accelerators and made the innovative design of the latter one of the most rapidly expanding fields of the AI industry. Various approaches have been explored to enable efficient and accurate processing of LLMs on the available accelerators given their computational and storage limitations. Among these, various quantization techniques have become the main focus of the community as a means of reducing the compute, communication and storage requirements. Quantization to lower precision formats naturally poses a number of challenges caused by the limited range of the available value representations. When it comes to processing the popular Transformer models on hardware, one of the main issues becomes calculation of the LayerNorm simply because accumulation of the variance requires a much wider dynamic range than the hardware enables. In this article, we address this matter and propose a computationally-efficient scaling technique that can be easily applied to Transformer models during inference. Our method suggests a straightforward way of scaling the LayerNorm inputs based on the static weights of the immediately preceding linear layers. The scaling factors are computed offline, based solely on the linear layer weights, hence no latency or computational overhead is added during inference. Most importantly, our technique ensures that no numerical issues such as overflow or underflow could happen during the compute. This approach offers smooth, accurate and resource-effective inference across a wide range of hardware architectures. The article provides theoretical justification as well as supporting numerical simulations.

Figures (5)

• Figure 1: Post-LN
• Figure 2: Pre-LN
• Figure 3: Residual branching options.
• Figure 4: The compute flow between consecutive LayerNorms of various Transformers. Pink blocks with capital letters stand for linear layers with the corresponding weight matrices, green $\mathcal{F}$-blocks represent non-linearities and $\mathbf{S}$-block represents $\operatorname{Softmax}$.
• Figure 5: Empirical histograms of the sum of squares in RMSNorm layers of the 9th decoder in Llama-2-7b, calculated on wikitext2. The red vertical cut-off line sets the maximal representable FP16 value (65k) beyond which FP16 overflows, the blue line shows the minimal normal FP16 value. Histograms (b) and (d) show that after SLaNC scaling no overflow (or underflow) is detected and the RMSNorm is computed precisely.