Towards Fully FP8 GEMM LLM Training at Scale

TL;DR

This paper presents a novel LLM training architecture that fully utilizes FP8 GEMMs during training, achieving high throughput and stability comparable to BF16, enabling scalable and efficient large language model pre-training.

Contribution

It introduces a new architecture supporting FP8 GEMMs throughout transformer training, addressing stability issues and enabling scalable, high-throughput LLM training at scale.

Findings

Supports FP8 GEMMs in all transformer GEMMs during training

Achieves throughput gains comparable to BF16 training

Maintains stable training with reduced outlier activations

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.