COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training (2410.19313v3)

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 LLM pretraining and fine-tuning and Vision LLM 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.

• The paper presents COAT, which utilizes dynamic range expansion to reduce quantization error by approximately 1.63× for optimizer states.
• COAT employs mixed-granularity activation quantization to achieve up to 1.65× memory reduction and 1.57× speedup in training activations.
• Experimental validation confirms that COAT enables larger batch sizes and full-parameter training on fewer GPUs while maintaining near-lossless model performance.

Memory-Efficient Model Training with FP8 Precision

The paper presents COAT (Compressing Optimizer States and Activations for FP8 Training), an innovative framework aimed at enhancing the memory efficiency of FP8 training for large-scale models. It leverages two primary techniques: Dynamic Range Expansion (DRE) and Mixed-Granularity Activation Quantization (MGAQ), addressing the limitations of existing FP8 training methodologies.

Key Contributions

1. Dynamic Range Expansion:
   • The authors introduce DRE to optimize the quantization of optimizer states by aligning their distributions with the FP8 representation range. This technique significantly reduces quantization error by adjusting the dynamic range effectively.
   • DRE achieves a quantization error reduction of approximately 1.63×, evidencing its capability to enhance optimizer state quantization accuracy.
2. Mixed-Granularity Activation Quantization:
   • COAT implements MGAQ to manage the quantization of activations, employing a hybrid approach with per-tensor and fine-grained quantization. This ensures a balance between accuracy and efficiency, particularly in non-linear layers.
   • It yields a memory reduction of 1.65× for activation storage compared to BF16, while the speedup is reported to reach 1.57× in specific scenarios.

Experimental Validation

The paper rigorously validates COAT across various tasks, such as LLM pretraining, fine-tuning, and VLM training. Key findings include:

• Efficiency Gains: COAT reduces memory usage by 1.54× compared to BF16, allowing full-parameter training on fewer GPUs and enabling batch size doubling in distributed training.
• Performance Stability: Despite extensive memory compression, the model maintains near-lossless performance across diverse datasets, demonstrating COAT's robustness.

Implications

The practical implications of COAT are profound, especially for researchers and practitioners facing hardware limitations. By enabling larger batch sizes and reducing memory consumption, COAT facilitates the training of expansive models with limited resources. Theoretically, this approach might inspire further exploration of low-precision training techniques and their applications to other neural network components.

Future Directions

Future research could explore:

• Integration with Communication-Efficient Techniques: Combining COAT with existing gradient compression methodologies might further optimize overall training efficiency.
• Adaptation to Different Architectures: Investigating the adaptability of COAT to different model architectures beyond those tested in the paper could widen its applicability.

Overall, COAT presents a promising approach to large-scale neural network training, addressing critical memory bottlenecks while preserving computational efficiency and model accuracy.