[Paper Review] Training and inference of large language models using 8-bit floating point

Sergio P. Perez, Yan Zhang+7 more|arXiv (Cornell University)|Sep 29, 2023

Parallel Computing and Optimization Techniques4 citations

TL;DR

This paper presents a scalable methodology for training and inferring large language models (up to 70B parameters) using 8-bit floating-point (FP8) formats by dynamically adjusting per-tensor scaling factors for weights, activations, and gradients. The FP8-AMAX approach achieves near-identical accuracy to full-precision FP16 baselines across GPT and Llama 2 models, demonstrating robustness even at scale.

ABSTRACT

FP8 formats are gaining popularity to boost the computational efficiency for training and inference of large deep learning models. Their main challenge is that a careful choice of scaling is needed to prevent degradation due to the reduced dynamic range compared to higher-precision formats. Although there exists ample literature about selecting such scalings for INT formats, this critical aspect has yet to be addressed for FP8. This paper presents a methodology to select the scalings for FP8 linear layers, based on dynamically updating per-tensor scales for the weights, gradients and activations. We apply this methodology to train and validate large language models of the type of GPT and Llama 2 using FP8, for model sizes ranging from 111M to 70B. To facilitate the understanding of the FP8 dynamics, our results are accompanied by plots of the per-tensor scale distribution for weights, activations and gradients during both training and inference.

Motivation & Objective

• Address the challenge of limited dynamic range in FP8 formats during training and inference of large language models.
• Overcome numerical instability from underflow and overflow in FP8 by dynamically adjusting scaling factors.
• Enable end-to-end FP8 training and inference for models up to 70B parameters with minimal accuracy degradation.
• Provide a practical, hardware-efficient alternative to FP16 and INT8 by leveraging FP8's computational advantages.
• Establish a reliable method for FP8 quantization that outperforms static scaling strategies in large models.

Proposed method

• Introduce FP8-AMAX, a dynamic scaling strategy that computes per-tensor scaling biases for weights, activations, and gradients based on the maximum absolute value of each tensor.
• Use a dynamic exponent bias adjustment via the compute_bias function, which calculates scaling based on log2 of the ratio between format maximum and tensor amplitude, with a configurable margin.
• Apply scaling before and after FP8 operations: scale inputs before casting to FP8, and unscale after matrix multiplication to preserve precision.
• Employ separate FP8 formats: FP8_E4 for weights and activations, and FP8_E5 for gradients, to balance range and precision.
• Integrate the method into both training (with gradient scaling and weight update) and inference (post-training quantization) pipelines.
• Use a loss scaling sweep to validate convergence robustness, comparing FP8-CSCALE (fixed scaling) against FP8-AMAX (adaptive scaling).

(a) Forward pass for FP16 inference.

Experimental results

Research questions

• RQ1 Can FP8 be effectively used for both training and inference of large language models without significant accuracy degradation?
• RQ2 How does dynamic per-tensor scaling (FP8-AMAX) compare to fixed scaling (FP8-CSCALE) in terms of convergence and stability across model sizes?
• RQ3 What scaling strategy ensures numerical stability in FP8 by preventing underflow and overflow in weights, activations, and gradients?
• RQ4 How do scaling biases evolve during training, and what implications does this have for update frequency and computational cost?
• RQ5 Can FP8 quantization achieve performance comparable to FP16 across diverse LLM architectures and tasks?

Key findings

• FP8-AMAX enables successful fine-tuning of GPT and Llama 2 models from 111M to 13B parameters with validation accuracy within 99.5% of the FP16 baseline across MNLI, QQP, and SST-2 tasks.
• For the 13B model, FP8-CSCALE only converges within a narrow scaling range of [-1, 0], and even then, convergence is not guaranteed across different random seeds.
• FP8-AMAX achieves consistent convergence across all model sizes tested, including 6.7B and 13B, with no divergence observed during fine-tuning.
• The loss function trajectories for FP8-AMAX and FP16 are nearly indistinguishable, indicating stable training dynamics.
• Per-tensor scaling distributions show that FP8-AMAX maintains stable, adaptive scaling throughout training, unlike the rigid scaling of FP8-CSCALE.
• The method successfully scales to 70B-parameter Llama 2 models, demonstrating feasibility of full FP8 training and inference at scale.

(b) Forward pass for FP8 inference.