PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling FP8 training to trillion-token LLMs

Maxim Fishman, Brian Chmiel, Ron Banner, Daniel Soudry

arXiv:2409.12517v2cs.LGcs.AI

TL;DR

Prior FP8 LLM training studies reached 100 billion tokens, leaving scalability and stability at larger scales unresolved. This work extends training to 2 trillion tokens, identifies SwiGLU-related outlier amplification, and introduces Smooth-SwiGLU alongside FP8 quantization of both Adam moments. The combined approach matches the BF16 baseline with up to 34% throughput improvement on a 7B model.

  • Problem

    Prior FP8 LLM studies reached 100 billion tokens, leaving scalability and stability at truly large training scales unresolved.

    Source: §1 · p. 1

  • Method

    The paper introduces Smooth-SwiGLU to mitigate outlier amplification and quantizes both Adam optimizer moments to FP8.

    Source: §1 · p. 1, §7 · p. 10

  • Results

    Up to 34% throughput improvement over the BF16 baseline was achieved with a 7B model trained using FP8 on 256 Intel Gaudi2 accelerators.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The combined method overcomes divergence challenges typically encountered in standard FP8 training on large datasets.

    Source: §7 · p. 10

    The weight-correlation analysis assumes an embedded SwiGLU neuron in a network trained with ℓ2 regularization.

    Source: §4.2 · p. 4

Abstract

from arXiv · show

We train, for the first time, large language models using FP8 precision on datasets up to 2 trillion tokens -- a 20-fold increase over previous limits. Through these extended training runs, we uncover critical instabilities in FP8 training that were not observable in earlier works with shorter durations. We trace these instabilities to outlier amplification by the SwiGLU activation function. Interestingly, we show, both analytically and empirically, that this amplification happens only over prolonged training periods, and link it to a SwiGLU weight alignment process. To address this newly identified issue, we introduce Smooth-SwiGLU, a novel modification that ensures stable FP8 training without altering function behavior. We also demonstrate, for the first time, FP8 quantization of both Adam optimizer moments. Combining these innovations, we successfully train a 7B parameter model using FP8 precision on 256 Intel Gaudi2 accelerators, achieving on-par results with the BF16 baseline while delivering up to a $\sim 34 \%$ throughput improvement. A reference implementation is supplied in https://github.com/Anonymous1252022/Megatron-DeepSpeed.

1 INTRODUCTION

The introduction presents extended FP8 LLM training as a way to address training costs, while showing that longer runs expose instabilities and describing two techniques for stable, efficient training.

1 INTRODUCTION

  • Earlier FP8 LLM studies reached datasets of up to 100 billion tokens, leaving scalability and stability at larger training scales open.The introduction frames lower-precision training as a response to the growing computational demands of LLMs.Source: §1 · p. 1
  • FP8 training on datasets up to 2 trillion tokens—20 times previous limits—revealed instabilities that shorter-duration studies did not observe.The introduction reports dramatic loss divergence after 220 billion tokens in an FP8 Llama2 7B run.Source: §1 · pp. 1–2
  • Smooth-SwiGLU mitigates outlier amplification while preserving the activation function’s behavior, supporting stable FP8 training over extended durations.The method is introduced as a modification to standard SwiGLU that does not alter model behavior.Source: §1 · pp. 1–2
  • The paper demonstrates FP8 quantization of both Adam optimizer moments for the first time, further reducing training memory use.The introduction presents this as an additional efficiency improvement for large-scale LLM development.Source: §1 · p. 2
  • Up to 34% throughput improvement over the BF16 baseline was achieved while training a 7B model on 256 Intel Gaudi2 accelerators using FP8.The combined approach achieved results on par with the BF16 baseline.Source: §1 · p. 2
  • The newly identified instability is outlier amplification by SwiGLU, which becomes pronounced only after extended training.The paper links amplification analytically and empirically to weight alignment in SwiGLUs with sufficiently large inputs during ℓ2-regularized training.Source: §1 · pp. 1–2

2 BACKGROUND AND CHALLENGES OF FP8 TRAINING IN LLMS

This section introduces FP8’s formats and dynamic-range challenge, reviews scaling approaches and their vulnerabilities, and positions the work as an extension to longer training runs.

2 BACKGROUND AND CHALLENGES OF FP8 TRAINING IN LLMS

  • FP8 uses E4M3 for weights and activations and E5M2 for gradients, following standardization for deep learning (Micikevicius et al., 2022).Earlier training formats include FP16 and BF16 (Micikevicius et al., 2017; Scao et al., 2022; Smith et al., 2022).Source: §2 · p. 2
  • FP8’s limited dynamic range makes careful scaling necessary to maintain numerical stability and model performance, including on H100 and Gaudi2 hardware.Source: §2 · p. 2
  • Global loss scaling prevents gradient underflow, while per-tensor scaling adapts to individual tensor ranges; instability can persist despite such techniques (Lee et al., 2024).Source: §2 · p. 2
  • Just-in-time scaling uses current data distributions but may require multiple passes, whereas delayed scaling uses preceding iterations and is vulnerable to outliers.Delayed scaling assumes statistical consistency across iterations.Source: §2 · p. 2
  • Building on FP8 LLM training up to 100 billion tokens (Peng et al., 2023), this work reaches 2 trillion tokens and introduces techniques for extended-run instability.The paper also reports improvements in memory use and training speed.Source: §2 · pp. 2–3

3 OUTLIER AMPLIFICATION IN LARGE-SCALE FP8 TRAINING

The section describes late-emerging activation outliers in large-scale training and explains why their sporadic appearance challenges FP8 stability.

3 OUTLIER AMPLIFICATION IN LARGE-SCALE FP8 TRAINING

  • After approximately 200 billion training tokens, activation maxima show sporadic outliers that challenge numerical stability and performance.The figure compares maxima across layers over 50 iterations and uses a changed z-axis scale for its late-training panel.Source: §3 · p. 3, Figure 1 · p. 3
  • Outliers exceeding FP8’s narrower representable range can cause overflow or underflow and destabilize training.The passage contrasts FP8’s range with FP32 and BF16.Source: §3 · p. 3
  • Because the outliers are sporadic, scaling methods that assume consistent statistics across iterations may struggle to adapt and training may become unstable or diverge.The section notes that these outliers become particularly prominent in later training stages.Source: §3 · p. 3

4 SWIGLU AND OUTLIER AMPLIFICATION

The section explains how SwiGLU can amplify outliers and why FP8 instability emerges over prolonged training, then presents Smooth-SwiGLU as a stable, function-preserving adjustment.

4.1 SWIGLU STRUCTURE

  • Unlike ReLU, GeLU, or Swish, SwiGLU grows quadratically for large inputs when w1 and w2 align, allowing unusually large outputs and outliers.The section identifies SwiGLU as a contributor to outlier amplification.Source: §4 · p. 3, §4.1 · p. 4

4.2 THEORETICAL ANALYSIS OF WEIGHT CORRELATION IN SWIGLU

  • At stationary points where SwiGLU input derivatives are small, the theory implies w1 and w2 align or anti-align; ℓ2 regularization drives alignment, and weak regularization can enlarge norms.The result also holds for other GLU variants; weak regularization strengthens the effect by increasing weight norms.Source: §4.2 · pp. 4–5

4.3 OBSERVING WEIGHT CORRELATION GROWTH AND TRAINING INSTABILITY

  • Correlation and outlier-channel norms rise sharply between 125B and 210B tokens; the resulting activation spikes lead to FP8 loss divergence after ∼200B tokens.Figure 2 tracks the loss, norms, correlation, weight elements, and outlier-channel histogram.Source: Figure 2 · p. 5, §4.3 · p. 6

4.4 SMOOTH-SWIGLU

  • Smooth-SwiGLU scales the SwiGLU linear branch and rescales after the last MLP linear layer, preventing quantized-input outliers while preserving the activation’s function and enabling full FP8.Figure 4 contrasts the original quantized SwiGLU with the proposed modification.Source: §4.4 · pp. 6–7, Figure 4 · p. 7

5 FP8 OPTIMIZER

The paper quantizes both Adam moments to FP8, reducing optimizer memory use. The two moments require different FP8 formats, with the second moment’s inverse-square-root role driving a wider-range choice.

5.1 CHALLENGES

  • The inverse square root of the second moment makes its smallest values most significant in parameter updates, creating a precision-reduction challenge.The first moment estimates the mean of gradients, while the second estimates their uncentered variance.Source: §5.1 · p. 8

5.2 METHODOLOGY

  • E4M3 for the first moment and E5M2 for the second is the only tested FP8 combination reported to converge to baseline.The first moment estimates the gradient mean; the second estimates uncentered gradient variance.Source: §5.2 · p. 8, Figure 5 · p. 8
  • The authors report the first successful quantization of both Adam moments using standard FP8 formats.Their experiments compare formats for each moment and find distinct precision requirements.Source: §5 · p. 7, §5.2 · p. 8

6 EXPERIMENTS

Experiments train Llama2 on Red Pajama for 2 trillion tokens using Intel Gaudi2 devices. The proposed configuration converges and achieves BF16-comparable results with reported acceleration and memory savings.

6.1 EXPERIMENTAL SETUP

  • The experiments use Llama2 trained on Red Pajama for 2 trillion tokens, with hyperparameters consistent with Touvron et al. (2023), on 256 Intel Gaudi2 devices.The training described for the proposed configuration ran for 15 days.Source: §6.1 · p. 9, §6.2 · p. 9

6.2 RESULTS

  • With Smooth-SwiGLU and both Adam moments quantized to FP8, Llama2 7B converges similarly to BF16, while standard FP8 training diverges after 200B tokens.The standard FP8 comparison uses SwiGLU and FP32 Adam moments.Source: Figure 6 · p. 9, §6.2 · p. 9
  • The proposed FP8 model achieves on-par zero-shot accuracy and perplexity with the BF16 baseline across all tested metrics.Table 2 compares the BF16 baseline with FP8 configurations, including Smooth-SwiGLU plus FP8 optimizer moments.Source: §6.2 · p. 9, Table 2 · p. 10
  • ∼34% acceleration with results on par with the BF16 baseline is reported for Smooth-SwiGLU.Table 3 compares acceleration configurations on 8 Intel Gaudi2 devices; the proposed scheme is reported to converge.Source: §6.2 · p. 9, Table 3 · p. 10
  • Changing optimizer moments from FP32 to FP8, alongside reducing the master weight to FP16, reduces memory consumption by ∼30%.Table 4 reports measurements on 8 Intel Gaudi2 devices using DeepSpeed Zero-1.Source: §6.2 · p. 10, Table 4 · p. 10

7 CONCLUSIONS

The paper reports FP8 training on datasets up to 2 trillion tokens and identifies prolonged-run outlier amplification by SwiGLU as a source of instability. It presents Smooth-SwiGLU and FP8 quantization of both Adam moments alongside BF16-comparable downstream performance.

7 CONCLUSIONS

  • Training to 2 trillion tokens, versus the previous 100-billion-token limit, revealed outlier-related FP8 instabilities that emerge over time and are amplified by SwiGLU.The paper reports that this behavior was not observable in earlier, shorter-duration studies.Source: §7 · p. 10
  • Smooth-SwiGLU applies per-channel quantization to SwiGLU, reducing outlier amplification while preserving the function and leaving inference unaffected.The authors report a moderate effect on model performance during training.Source: §7 · p. 10
  • The paper introduces FP8 quantization of both Adam optimizer moments as a way to further optimize memory usage.The authors describe this as the first implementation of this approach.Source: §7 · p. 10
  • The combined method achieves downstream-task performance comparable to BF16 baselines while providing throughput improvements.The conclusion also states that it overcomes divergence challenges encountered in standard FP8 training on large datasets.Source: §7 · p. 10
  • The paper states that its anonymous GitHub repository contains code and details for reproducing the experiments.The supplied repository is identified in the paper abstract.Source: §7 · p. 11

A.1 TRAINING INSTABILITY - ADDITIONAL DATA

The appendix compares SwiGLU weight outliers early and late in training and reports that most values in a late-training outlier channel fall in a range where the activation derivative is very small.

A.1 TRAINING INSTABILITY - ADDITIONAL DATA

  • Outlier elements in w1 and w2 show minimal correlation at 8B tokens but high negative correlation at 330B tokens.The figure also contrasts histograms of an outlier channel at those two training stages.Source: Figure 7 · p. 15
  • At 200B tokens, ∼1% of values are below 1 and ∼3.5% below e; σ′(w⊤ 2 xn) is very small for the overwhelming majority of n values.The histogram uses natural-log scale on the x-axis and base-10 log scale on the y-axis.Source: Figure 9 · p. 15, §A.1 · p. 15

A.2 PERFORMANCE GAIN ON NVIDIA GPUS

The section extends the performance-acceleration comparison to Nvidia GPUs, using a non-optimized Llama2 7B implementation measured on 8 Nvidia A6000 Ada GPUs.

A.3 SMOOTH-SWIGLU STUDY

Smooth-SwiGLU produces smoother BF16 training and can reach lower loss, especially at larger learning rates; a separate GPT-3 FP8 run without SwiGLU showed no stability issues.

A.3 SMOOTH-SWIGLU STUDY

  • Smooth-SwiGLU makes BF16 training smoother and enables lower loss, especially at larger learning rates.The study examines different learning rates; 2.5e-4 is the standard baseline with a cosine scheduler.Source: §A.3 · p. 16, Figure 10 · p. 16, Figure 11 · p. 16

A.4 FP8 WITHOUT SWIGLU ACTIVATION FUNCTION

Loading 2409.12517v2…