COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training

TL;DR

<3-5 sentence high-level summary> COAT addresses memory bottlenecks in FP8 training by quantizing both optimizer states and activations to FP8, introducing Dynamic Range Expansion to better fit FP8’s range and a Mixed-Granularity Activation Quantization strategy to balance precision and efficiency. The method preserves accuracy across LLM pretraining/fine-tuning and VLM tasks while delivering significant end-to-end memory reductions (≈1.54x) and speedups (≈1.43x), enabling full-parameter training on fewer GPUs and larger batch sizes. This work demonstrates practical scalability of FP8 for large foundation models and provides a public implementation to facilitate adoption. Its significance lies in making memory-efficient full-parameter training feasible on commodity-scale GPU clusters for billion-parameter models.

Abstract

FP8 training has emerged as a promising method for improving training efficiency. Existing frameworks accelerate training by applying FP8 computation to linear layers while leaving optimizer states and activations in higher precision, which fails to fully optimize memory usage. This paper introduces COAT (Compressing Optimizer States and Activations for FP8 Training), a novel FP8 training framework designed to significantly reduce memory footprint when training large models. COAT addresses current limitations through two key innovations: (1) Dynamic Range Expansion, which aligns optimizer state distributions more closely with the FP8 representation range, thereby reducing quantization error, and (2) Mixed-Granularity Activation Quantization, which optimizes activation memory using a combination of per-tensor and per-group quantization strategies. Experiments demonstrate that COAT effectively reduces end-to-end training memory footprint by 1.54x compared to BF16 while achieving nearly lossless performance across various tasks, such as Large Language Model pretraining and fine-tuning and Vision Language Model training. COAT also achieves a 1.43x end-to-end training speedup compared to BF16, performing on par with or surpassing TransformerEngine's speedup. COAT enables efficient full-parameter training of large models on fewer GPUs, and facilitates doubling the batch size in distributed training settings, providing a practical solution for scaling large-scale model training. The code is available at https://github.com/NVlabs/COAT.

Figures (10)

• Figure 1: (a,b) Comparing the quantization flow of Transformer Engine and COAT. Both the optimizer states and activations are quantized to FP8 in COAT. (c) End-to-end per-GPU memory comparison when training Llama-2-13B on 8$\times80$G H100 using FSDP.
• Figure 1: Quantization error of $\frac{m}{\sqrt{v}}$ under different quantization settings. +Expand means applying our Dynamic Range Expansion method.
• Figure 2: (a) Visualization of optimizer states' dynamic range under per-group quantization. FP8 E4M3's representation range is under-utilized in this case. (b) After dynamic range expansion, FP8's representation range is well utilized. (c) Distribution of $k$ for optimizer states. The second order's $k$ is larger than the first order's $k$, since the second-order momentum's dynamic range is smaller.
• Figure 3: Dynamic Range Expansion can better utilize E4M3 representation range.
• Figure 4: (a) Quantization Error in forward pass. (b) Time comparison of various scaling methods.

Theorems & Definitions (1)

• Definition 1: dynamic range