PaperlayerView exact revision on arXiv

Source-linked AI summary

FP8-LM: Training FP8 Large Language Models

Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, Ruihang Li, Miaosen Zhang, Chen Li, Jia Ning, Ruizhe Wang, Zheng Zhang, Shuguang Liu, Joe Chau, Han Hu, Peng Cheng

arXiv:2310.18313v2cs.LGcs.CL

TL;DR

LLM training is costly, and existing FP8 support leaves some weights and gradients in high precision. This paper proposes a staged FP8 mixed-precision framework and reports reduced memory and communication costs versus BF16 alongside performance equivalency.

  • Problem

    LLM training is extremely costly, and existing FP8 support is limited to GEMM while retaining high-precision master weights and gradients.

    Source: §1 · pp. 1–2

  • Method

    The paper introduces a staged FP8 mixed-precision framework that incrementally applies 8-bit compute, storage, communication, optimizer, and distributed training, with precision decoupling and dynamic scaling.

    Source: §1 · pp. 2–3

  • Results

    29% to 39% reduction in real memory usage and 63% to 65% decrease in weight-related communication overhead versus BF16 accompanied performance equivalency without hyperparameter changes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework reduces costs in GPT training and extends to instruction tuning and RLHF, where FP8 reduces model-weight and optimizer-state memory.

    Source: §1 · pp. 1, 3

Abstract

from arXiv · show

In this paper, we explore FP8 low-bit data formats for efficient training of large language models (LLMs). Our key insight is that most variables, such as gradients and optimizer states, in LLM training can employ low-precision data formats without compromising model accuracy and requiring no changes to hyper-parameters. Specifically, we propose a new FP8 automatic mixed-precision framework for training LLMs. This framework offers three levels of FP8 utilization to streamline mixed-precision and distributed parallel training for LLMs. It gradually incorporates 8-bit gradients, optimizer states, and distributed learning in an incremental manner. Experiment results show that, during the training of GPT-175B model on H100 GPU platform, our FP8 mixed-precision training framework not only achieved a remarkable 39% reduction in real memory usage but also ran 75% faster than the widely adopted BF16 framework (i.e., Megatron-LM), surpassing the speed of Nvidia Transformer Engine by 37%. This largely reduces the training costs for large foundation models. Furthermore, our FP8 mixed-precision training methodology is generic. It can be seamlessly applied to other tasks such as LLM instruction tuning and reinforcement learning with human feedback, offering savings in fine-tuning expenses. Our FP8 low-precision training framework is open-sourced at {https://github.com/Azure/MS-AMP}{aka.ms/MS.AMP}.

1 Introduction

LLM training is costly, while existing FP8 support leaves much of training storage and communication in high precision. The paper proposes an incremental FP8 framework and reports reduced memory and communication costs alongside performance equivalency to BF16.

1 Introduction

  • FP8 theoretically offers 2× speed and 50% - 75% memory and communication savings over 16-bit and 32-bit mixed-precision training, but existing systems default to FP32 or FP16/BF16.The introduction motivates these savings by the high cost of training LLMs.Source: §1 · pp. 1–2
  • Precision decoupling assigns reduced precision to less-sensitive components, while dynamic tensor scaling keeps gradients within FP8 range to alleviate underflow and overflow during all-reduce.The paper identifies FP8’s narrower dynamic range and reduced precision as sources of numerical instability.Source: §1 · p. 2
  • FP8 training reduces real memory usage by 29% to 39% and weight-related communication overhead by 63% to 65% versus BF16 mixed-precision; without hyperparameter changes, it matches BF16 high-precision performance in pre-training.Source: §1 · p. 2
  • 37% less training time than Transformer Engine and 42% less memory on H100 were reported for GPT-175B training.The comparison is with Transformer Engine (TE).Source: §1 · p. 3
  • FP8 fine-tuning matched BF16 on AlpacaEval and MT-Bench while improving training speed by 27%.In RLHF, AlpacaFarm with FP8 reduced model-weight memory by 32% and optimizer-state memory by 62%.Source: §1 · p. 3
  • The framework incrementally adds 8-bit weights, gradients, optimizer, and distributed training as a drop-in replacement requiring no hyperparameter or training-recipe changes.The paper also describes GPT-style models from 7B to 175B parameters and FP8 support for tensor, pipeline, and sequence parallelism.Source: §1 · pp. 2–3

2 FP8 LLMs

The section develops an FP8 training strategy that addresses numerical instability while extending low-precision use from gradient communication and optimizer storage to distributed parallelism. Its components use automatic scaling, precision choices tailored to optimizer variables, and FP8-aware tensor distribution.

2 FP8 LLMs

  • FP8 training faces narrower dynamic range and lower representation precision than BF16 and FP16, raising the risk of loss spikes or NaNs.Tensor scaling multiplies higher-precision values before FP8 casting to better align them with the FP8 format's representable range.Source: §2 · pp. 3–4

2.1 FP8 Gradient and All-Reduce Communication

  • An automatic scaling factor changes during training to reduce gradient overflow and underflow, while one shared scalar lets all GPUs quantize tensors associated with the same weight into FP8.If the maximum-value ratio exceeds 0.001%, µ is set to 1/2 in the next training step; if it consistently remains at the threshold, µ is exponentially increased to 2 over 1,000 training steps.Source: §2.1 · pp. 4–5

2.2 FP8 Optimizer

  • The FP8 optimizer assigns lower precision to the first-order moment but higher precision to master weights and the second-order moment.The authors report that their optimizer preserves accuracy from 125M to 175B parameters and reduces optimizer memory footprints by 2.6x.Source: §2.2 · pp. 5–6

2.3 FP8 Distributed Parallel Training

  • FP8 tensor and sequence parallelism use low-bit activations for gather-reduce communication, while FP8-aware ZeRO distributes whole scaled tensors rather than partitions.FP8 conversion before the sequence/tensor-parallel gather-reduce operation saves communication; whole-tensor distribution lets scaling factors travel with tensors.Source: §2.3 · pp. 6–7, Figure 3 · p. 7

3 Experiment

The experiments evaluate FP8 training across GPT pre-training, fine-tuning, and system performance. FP8 maintains performance comparable to BF16 while reducing memory and communication costs and improving training throughput.

3 Experiment

  • Experiments cover GPT models from 125M to 175B parameters, pre-training and fine-tuning datasets, and H100 training; the 175B model uses 40B tokens for system evaluation.Pre-training data combines open-source collections; instruction tuning uses ShareGPT, while RLHF data combines Anthropic’s Helpful and Harmless and Open-Assistant datasets.Source: §3 · pp. 7–8, §3.1.1 · p. 8, §3.1.2 · p. 8

3.2.1 Model Performance

  • FP8 and BF16 training losses nearly overlap across GPT models from 7B to 175B, and their zero-shot performance is comparable across downstream tasks.Configurations and hyperparameters are consistent across the precision schemes; the only difference is mixed precision.Source: Figure 4 · p. 9, Table 2 · p. 9, §3.2.1 · p. 9
  • SFT results are comparable between FP8 and BF16, while RLHF with FP8 reduces memory use by 32% for weights and 62% for optimizer states.The SFT comparison reports GPU memory and throughput alongside AlpacaEval win-rate and MT-Bench scores.Source: §3.2.1 · pp. 9–10, Table 4 · p. 10, Figure 6 · p. 10

3.2.2 System Performance

  • Using 8-bit gradients for all-reduce theoretically cuts communication costs by 75% versus the 32-bit scheme; observed GPT reductions are 63% to 65%.The passage attributes the smaller practical reduction to system transmission loss.Source: §3.2.2 · p. 10, Table 5 · p. 11
  • 75% speedup on GPT-175B; FP8 reaches 34.2% MFU on H100 GPUs, 37.3% higher than TE.At identical batch sizes, memory footprint falls 28% to 39% versus BF16; FP8 also provides 36.1%, 36.0%, and 42.1% additional memory reductions versus TE on GPT-7B, 13B, and 175B.Source: §3.2.2 · pp. 10–11, Table 5 · p. 11

3.3 Ablation Study

  • Auto-scaling reduces both underflow and overflow and improves SNR over pre-scaling and post-scaling for FP8 gradient all-reduce.The comparison measures SNR, underflow, and overflow across Transformer blocks on GPT-7B with data parallelism factor 128.Source: §3.3 · p. 11, Figure 7 · p. 11
  • FP8 master weights degrade performance; scaled FP16 preserves FP32 accuracy, BF16 is slightly worse, and FP8 second moments can cause divergent loss through underflow.The optimizer ablation compares master weights and first- and second-order gradient moments.Source: §3.3 · p. 12, Table 6 · p. 12, Figure 8 · p. 12
  • FP8 parallelism reduces activation-related communication costs by 34%; distributing each FP8 tensor with its scaling factor supports lower memory use and balanced GPU loads.The ZeRO method distributes each tensor and its scaling factor as a whole rather than splitting the tensor across GPUs.Source: §3.3 · p. 12, Table 7 · p. 12, Table 8 · p. 12

4 Related Work

LLM training has moved from FP16 toward BF16 for stability, while FP8 became practical with Hopper hardware. Earlier FP8 work and Transformer Engine support remained limited, motivating broader FP8 use across training.

Mixed-precision training

  • FP16 mixed precision uses FP16 for forward and backward computations with FP32 master weights, while loss scaling addresses FP16’s narrower dynamic range (Micikevicius et al., 2017).BF16 avoids loss scaling by retaining FP32’s dynamic range.Source: §4 · p. 13

FP8 training

  • Earlier FP8 training work largely remained at simulation stage, and Transformer Engine restricts FP8 to weight computation while storing weights and gradients in 16-bit formats (Micikevicius et al., 2022).The paper presents FP8 gradient, optimizer, and distributed training as its broader alternative.Source: §4 · p. 13

Large language models

  • Autoregressive language modeling predicts future text from past text, and scaling-law work accompanied the development of dense and sparse LLMs (Brown et al., 2020; Hoffmann et al., 2022).These models still face substantial compute requirements and demands for high-quality training data.Source: §4 · p. 13

Mixed-precision training

  • BF16 became preferred for large-model training because it retains FP32’s dynamic range, whereas FP16 can cause instability, particularly beyond 100B parameters (Scao et al., 2022; Rae et al., 2021; Hoffmann et al., 2022).OPT and GLM use FP16 for forward and backward passes but retain FP32 optimizer states and master weights (Zhang et al., 2022; Zeng et al., 2022).Source: §4 · p. 13

5 Conclusion

The paper presents an incremental FP8 training framework for LLMs and reports reduced communication overhead and memory use in GPT training. It identifies larger-scale training, multimodal models, and low-bit deployment as future directions.

Conclusion

  • The framework incrementally incorporates 8-bit collective communication, optimizer, and distributed parallel training into LLM training.The authors describe this as extending FP8 compute, storage, and communication across the training process.Source: §5 · p. 14
  • Experiments on GPT models at various scales show reduced communication overhead and memory utilization.The authors plan to scale up model size and training steps in future work.Source: §5 · p. 14
  • Planned extensions include multimodal model training and low-bit LLM deployment on edge devices such as smartphones.Source: §5 · p. 14

Contribution and Acknowledgement

The acknowledgements identify project leadership and assign co-author contributions across framework, engineering, systems, hardware, data, training, alignment, and evaluation.

Acknowledgements

Contributions

Contributions and acknowledgements

Contributions

  • Product engineering, systems, and hardware infrastructure contributions are assigned to separate listed teams.The team members are listed in the engineering, system, and hardware entries.Source: Contribution and Acknowledgement · p. 14

A.1 FP8 Data Formats

FP8 includes E4M3 and E5M2 sub-formats that trade representation range against precision. Its narrower range and reduced precision create distinct challenges for large-model training.

A.1 FP8 Data Formats

  • The FP8 specification defines E5M2 and E4M3 as formats trading a larger range for higher precision, respectively.The specification was published by NVIDIA, ARM, and Intel in September 2022 as an AI interchange format.Source: §A.1 · p. 20
  • E4M3 stores values up to +/-448 and NaN, using 1 sign bit, 4 exponent bits, and 3 mantissa bits.Source: §A.1 · p. 20
  • E5M2 stores values up to +/-57344, +/- inf, and NaN, using 1 sign bit, 5 exponent bits, and 2 mantissa bits.Source: §A.1 · p. 20
  • FP8’s narrower representation range and lower representation precision are significant challenges for training large models.The paper attributes the range limit primarily to exponent bits and quantization error to the limited mantissa bits.Source: §A.1 · p. 20

A.2 FP8 Tensor Scaling

Tensor scaling brings values into FP8’s representable range, with fine-grained per-tensor scaling supporting stable training despite FP8’s narrower range. The scaling-factor strategies trade execution overhead against storing history from prior iterations.

A.2 FP8 Tensor Scaling

  • Global loss scaling can still encounter significant underflow in super-large models or complex tasks, prompting block-wise and layer-wise gradient scaling.Global scaling uses one adaptive factor for gradients across all layers.Source: §A.2 · p. 21
  • FP8 per-tensor scaling enables stable training within ranges of [1.95E-3, 448] for E4M3 and [1.53E-5, 5.73E+4] for E5M2.Figure 9 visualizes scaling tensor values into the datatype’s representation range.Source: §A.2 · p. 21, Figure 9 · p. 21
  • Just-in-time scaling requires multiple data passes and intermediate higher-precision output, while delayed scaling uses preceding iterations’ maxima but stores their history.Just-in-time overhead can substantially reduce FP8 benefits; delayed scaling permits full FP8 computation performance.Source: §A.2 · p. 21

A.3 Pre-training Data

The pre-training data section describes the sources and preprocessing of the collected corpora, while Table 10 summarizes sampling weights, epochs, and token counts. For runs under 100 billion training tokens, the same sampling proportions are used.

A.3 Pre-training Data

  • The sources include Redpajama’s arXiv and StackExchange subsets, The Pile subsets, and a 24-language Wikipedia dump dated 20220301.The Pile subsets include BookCorpus2, Books3, DM-Math, Gutenberg, HackerNews, NIH ExPorter, OpenSubtitles, and USPTO.Source: §A.3 · p. 22
  • Eleven CommonCrawl snapshots from 2018–2023 undergo line-level and fuzzy deduplication, English-language filtering, and quality and Wikipedia-similarity filtering.The pipeline uses CCNet, a fastText language classifier, an n-gram language model, and a linear classifier.Source: §A.3 · p. 22
  • Python code is collected from GitHub and Stack, filtered by content and file characteristics, then fuzzy-deduplicated.Filters include line and file alphanumeric rates, file length, maximum and average line length, and Python keyword frequency.Source: §A.3 · p. 22
  • Table 10 lists each pre-training subset’s sampling weight, number of epochs, and training tokens, with the same sampling proportions used below 100 billion training tokens.Books data combines BookCorpus2, Books3, and Gutenberg; dialogue data combines HackerNews and OpenSubtitles.Source: Table 10 · p. 23
Loading 2310.18313v2…