Towards Fully FP8 GEMM LLM Training at Scale

Alejandro Hernández-Cano, Dhia Garbaya, Imanol Schlag, Martin Jaggi

Keywords: FP8, LLMs, Training Dynamics, Quantization, Efficient Training

TL;DR: We perform an important step towards LLM pure FP8 training by enabling stable FP8 dot product attention reaching new throughput records

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. Additionally, we identify key metrics for monitoring low-precision training and predicting potential future divergences.

Primary Area: Deep learning (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 22812