COAT : Compression des états de l'optimiseur et activation pour un entraînement efficace en FP8.

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

October 25, 2024

Auteurs: Haocheng Xi, Han Cai, Ligeng Zhu, Yao Lu, Kurt Keutzer, Jianfei Chen, Song Han

cs.AI

Résumé

La formation en FP8 s'est imposée comme une méthode prometteuse pour améliorer l'efficacité de l'entraînement. Les cadres existants accélèrent l'entraînement en appliquant le calcul en FP8 aux couches linéaires tout en laissant les états de l'optimiseur et les activations dans une précision plus élevée, ce qui ne permet pas d'optimiser pleinement l'utilisation de la mémoire. Cet article présente COAT (Compression des États de l'Optimiseur et des Activations pour la Formation en FP8), un nouveau cadre de formation en FP8 conçu pour réduire de manière significative l'empreinte mémoire lors de l'entraînement de grands modèles. COAT aborde les limitations actuelles à travers deux innovations clés : (1) l'Expansion de la Plage Dynamique, qui aligne plus étroitement les distributions des états de l'optimiseur avec la plage de représentation en FP8, réduisant ainsi l'erreur de quantification, et (2) la Quantification des Activations à Granularité Mixte, qui optimise la mémoire d'activation en utilisant une combinaison de stratégies de quantification par tenseur et par groupe. Les expériences montrent que COAT réduit efficacement l'empreinte mémoire d'entraînement de bout en bout de 1,54x par rapport au BF16 tout en atteignant des performances presque sans perte sur diverses tâches, telles que la préformation et le fine-tuning de grands modèles de langage et la formation de modèles de langage visuel. COAT réalise également un gain de vitesse d'entraînement de bout en bout de 1,43x par rapport au BF16, se situant au niveau ou dépassant le gain de vitesse de TransformerEngine. COAT permet un entraînement efficace de tous les paramètres de grands modèles sur moins de GPU, et facilite le doublement de la taille du lot dans des paramètres d'entraînement distribué, offrant une solution pratique pour mettre à l'échelle l'entraînement de modèles à grande échelle. Le code est disponible sur https://github.com/NVlabs/COAT.

English

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.