Research map / Architectures and efficiency / RNNs and normalization
Towards Fully FP8 GEMM LLM Training at Scale
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces FOG, a family of outlier-guarded transformer architectures that enables FP8 GEMMs for all linear and attention matmuls in transformer blocks (FP8DPA) with delayed scaling, matching BF16 quality while training up to 43% faster at 8B scale.
Problem
FP8 LLM pre-training is unstable at scale, and existing recipes keep attention GEMMs in higher precision or use costly fine-grained scaling, limiting throughput gains.
Approach
Redesign the transformer block (post-normalization instead of pre-normalization, frozen or tanh-based QK normalization, scaled softmax, and non-gated activations such as xIELU or GeLU) so that all GEMMs inside transformer blocks, including attention QK^T and PV, can run in FP8 with low-overhead tensorwise delayed scaling; also uses kurtosis of activations to predict late divergence.
Key results
- FOG-flash FP8DPA reaches 12764 tokens/s/GPU at 8B (+40.2% over BF16 Llama3), with 42.6% in the Zero-1 setting reported elsewhere.
- FOG-max FP8DPA at 1.5B matches BF16 Llama3 average downstream accuracy (46.0 BF16 vs 47.1 FP8DPA reported for FOG-max), while OP, OLMo2 and Llama3 variants diverge in FP8DPA training.
- Llama3 divergence marks occur earlier for smaller models (0.7B tokens at 390M vs 6.6B tokens at 8B), and a 1.5B FOG-max run trained stably for 450B tokens.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use FOG-flash with FP8 GEMMs in attention | 8B LLM training · context 4096 | tokens/second/GPU: 12764 (+40.2%) | Helped |
| Use FOG-flash with FP8 GEMMs in attention | 1.5B LLM training · context 4096 | tokens/second/GPU: 54848 (+18.0%) | Helped |
| Use FOG-flash FP8DPA instead of Llama+SmoothSwiGLU FP8 | 8B MoE training · MoE | tokens/second/GPU: 4351 (+29.2%) vs 4202 (+24.8%) | Helped |
| Use FP8 optimizer moments (ablation) | FOG-flash FP8DPA · 390M-token setup | final loss: 2.649 -> 2.645 | No clear effect |
| Freeze QK RMSNorm gains (ablation) | OP under FP8DPA · training stability | divergence | Helped |
| Add post-normalization (ablation) | OP with frozen QK gains under FP8DPA · learning-rate decay phase | convergence | Helped |
| Replace LayerScale-only residual normalization with post-normalization (ablation) | OP under FP8DPA · learning-rate decay phase | convergence | Helped |
| Cool down weight decay during learning-rate decay (ablation) | FOG-opt · 1.5B training | Average score: 46.0 ±0.3 -> 46.3 ±0.3 | No clear effect |
| Cool down weight decay during learning-rate decay (ablation) | OP+FrozenQK · learning-rate decay phase | divergence | No clear effect |
| Replace FOG-max xIELU with SwiGLU (ablation) | FOG-max · FP8DPA training | stability | No clear effect |
| Replace FOG-max xIELU with SwiGLU (ablation) | FOG-max · 1.5B FP8DPA training | stability | No clear effect |
| Use FP8DPA instead of BF16 | FOG-flash MoE · 1.8B MoE | final loss: 2.477 -> 2.483 | Hurt |
Surprising findings
- Larger models diverge later in FP8 training than smaller ones across architectures (e.g., OLMo2 at 390M diverges at 3.3B tokens vs 15.9B at 1.5B).
- A quadratic activation (xIELU) remained stable under FP8 attention when kurtosis was low, challenging prior explanations that attributed late Llama divergence solely to gated quadratic activations.
Limitations
- The LM head remains in BF16, so FP8 is not applied to all GEMMs.
- Experiments are limited to models up to 8B parameters and 20B tokens at 8B scale.
- Only tensorwise delayed scaling was evaluated; fine-grained scaling was compared only briefly.
Experimental setup
Datasets: FineWeb-Edu
Benchmarks: Hellaswag ARC-Easy PIQA CommonsenseQA LAMBADA-OpenAI LAMBADA-standard OpenBookQA SocialIQA WinoGrande
Metrics: accuracy training throughput (tokens/second/GPU) training loss kurtosis
Baselines: Llama3 BF16 Llama3+SmoothSwiGLU FP8 OP OLMo2 Llama3 FP8 blockwise scaling
Builds on: Outlier Protected (OP) architecture SmoothSwiGLU Megatron-LM Transformer Engine FP8 recipes Llama3
Closest papers by meaning
- Scaling FP8 training to trillion-token LLMs (ICLR 2025 · similarity 0.91)
- µnit Scaling: Simple and Scalable FP8 LLM Training (ICML 2025 · similarity 0.91)
- Optimizing Large Language Model Training Using FP4 Quantization (ICML 2025 · similarity 0.90)
- Quartet: Native FP4 Training Can Be Optimal for Large Language Models (NeurIPS 2025 · similarity 0.89)
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.88)
- Learning from Students: Applying t-Distributions to Explore Accurate and Efficient Formats for LLMs (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.