Towards Fully FP8 GEMM LLM Training at Scale

TL;DR

This work addresses FP8-based training for large language models at scale, proposing the FOG architecture family to enable FP8 GEMMs across all transformer block operations, including attention. By removing pre-normalization, introducing attention-time normalization, and applying post-normalization with careful activation choices, FOG achieves stable FP8DPA training and up to 43% throughput gains at 8B scale while maintaining BF16-level downstream performance. The paper introduces kurtosis as a diagnostic to monitor long-term outlier dynamics and to predict potential divergences, enabling more reliable FP8 training across long data regimes and MoE variants. Empirically, FOG variants converge where prior FP8 approaches diverge, demonstrate stable long-data training up to 450B tokens, and exhibit robustness across activation functions and MoE configurations, marking a significant step toward fully FP8 GEMM training in modern LLMs.

Abstract

Despite the significant potential of FP8 data formats for large language model (LLM) pre-training, their adoption has been limited due to challenges in maintaining stability at scale. Existing approaches often rely on suboptimal fine-grained FP8 kernels or fall back to higher-precision matrix multiplications (GEMMs) in sensitive components, such as attention projections, compromising potential throughput gains. We introduce a new class of LLM architectures that, for the first time, support FP8 computation for all GEMMs within transformer blocks during both forward and backward passes. This enables unprecedented throughput gains, particularly at scale, while matching the downstream performance of standard BF16 training. Our architecture design reduces large outlier activations, promoting stable long-term FP8 training. In addition, we identify key metrics to monitor low-precision training and predict potential future divergences.

Figures (10)

• Figure 1: FOG transformer.
• Figure 2: From OP to FOG-opt step by step. Comparison of 390M models under FP8DPA training. The first architecture to diverge is OP, while OP with frozen QK RMSNorm gains survives the stable phase of training. It still, however, experiences a significant divergence during the learning rate cooldown, which starts around 42B tokens in. The architecture that converges, FOG-opt, is the result of adding post-normalization to the previous recipe. Gradient norm reported is the 200-rolling-window mean and 5%-95% quantile bands.
• Figure 3: Kurtosis of QKV tensors during FP8DPA learning rate cooldown with OP+frozenQK architecture. Later layers exhibit significantly larger activation outliers.
• Figure 4: Loss and kurtosis training dynamics of 1.5B FOG-max and Llama3 models trained for over 100B tokens with BF16 precision. Loss reported is the 200-rolling-window mean and 5%-95% quantile bands.
• Figure 5: Training dynamics of a failed and a successful FP8DPA run. Kurtosis exhibits atypical behaviour much earlier than when the loss diverged. Gradient norm reported is the 200-rolling-window mean and 5%-95% quantile bands.

Theorems & Definitions (1)

• Theorem 1