To FP8 and Back Again: Quantifying Reduced Precision Effects on LLM Training Stability

TL;DR

The paper tackles the challenge of making FP8 feasible for LLM pretraining by quantifying how reduced-precision bit-width affects training stability relative to BF16. It introduces a novel loss-landscape sharpness metric and uses incremental bit-reduction experiments with masked GEMMs to probe stability, revealing that current FP8 methods (including MS-AMP) are not yet robust enough for cost-effective deployment. Key findings show exponent-bit reductions are the main destabilizers, and stability narrows the hyperparameter space, often depending on data quality and architectural choices. The work calls for rigorous, model- and data-agnostic stability analyses and suggests stabilization strategies like staged precision and selective high-precision layers to realize potential FP8 gains in real-world LLM training.

Abstract

The massive computational costs associated with large language model (LLM) pretraining have spurred great interest in reduced-precision floating-point representations to accelerate the process. As a result, the BrainFloat16 (BF16) precision has become the de facto standard for LLM training, with hardware support included in recent generations of accelerators. This trend has gone even further in the latest processors, where FP8 has recently been introduced. However, prior experience with FP16, which was found to be less stable than BF16, raises concerns as to whether FP8, with even fewer bits than FP16, can be a cost-effective option for LLM training. We argue that reduced-precision training schemes must have similar training stability and hyperparameter sensitivities to their higher-precision counterparts in order to be cost-effective. However, we find that currently available methods for FP8 training are not robust enough to allow their use as economical replacements. This prompts us to investigate the stability of reduced-precision LLM training in terms of robustness across random seeds, learning rates, and datasets. To this end, we propose new evaluation techniques and a new metric for quantifying loss landscape sharpness in autoregressive language models. By simulating incremental bit reductions in floating-point representations, we analyze the relationship between representational power and training stability with the intent of aiding future research into the field.

Figures (12)

• Figure 1: We show an example case of loss divergence on nanoGPT when using the same configurations as the default run except for the random seed. The blue lines indicate the average losses obtained for eight training runs that did not diverge. Of the 188 random seeds that were tested, 18 were found to diverge. As full pretraining requires over 4 days on a single node with 8 A100 GPUs, even for BF16, we perform early stopping at 30K steps, or 5% of the original training steps, requiring approximately 4 hours for a BF16 run and 8 hours for a TF32 run per A100 node with 8 GPUs. Because we only run 5% of the original training, we suspect that the measured divergence rate of approximately 10% underestimates the true rate of training loss divergence.
• Figure 2: Loss landscape diagrams for Llama 120M E8M3 at 5K steps (left) and 10K steps (right). Even during loss divergence, the loss landscape visualized using the method in visualloss appears smooth, motivating our introduction of a new loss landscape sharpness metric. The validation loss is shown for each point of the loss landscape.
• Figure 3: PyTorch-like pseudocode for the forward pass.
• Figure 4: Diagram showing the precisions used in a Llama decoder block (best seen in color). The activations in the path of the residual connection are kept in FP32, as are the model weights and embeddings. The LayerNorm and RoPE RoFormer layers use FP32 internally for their computations. The Flash Attention kernel uses BF16 with no reduction in precision. All other layers use reduced-precision matrix multiplication that emulates low-precision computation with a high-precision accumulator.
• Figure 5: Comparison between MS-AMP FP8 (O1) and BF16 training on a subsample of the FineWeb Edu dataset using a Llama 120M model.