COAT: Compresión de estados de optimizador y activación para un entrenamiento eficiente en FP8.

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

October 25, 2024

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

cs.AI

Resumen

El entrenamiento en FP8 ha surgido como un método prometedor para mejorar la eficiencia del entrenamiento. Los marcos existentes aceleran el entrenamiento aplicando cálculos en FP8 a capas lineales mientras mantienen los estados del optimizador y las activaciones en una precisión superior, lo que no logra optimizar completamente el uso de la memoria. Este documento presenta COAT (Compresión de Estados del Optimizador y Activaciones para el Entrenamiento en FP8), un marco de entrenamiento en FP8 novedoso diseñado para reducir significativamente la huella de memoria al entrenar modelos grandes. COAT aborda las limitaciones actuales a través de dos innovaciones clave: (1) Expansión del Rango Dinámico, que alinea las distribuciones de estados del optimizador de manera más cercana con el rango de representación en FP8, reduciendo así el error de cuantificación, y (2) Cuantificación de Activaciones de Granularidad Mixta, que optimiza la memoria de activación utilizando una combinación de estrategias de cuantificación por tensor y por grupo. Los experimentos demuestran que COAT reduce efectivamente la huella de memoria de entrenamiento de extremo a extremo en 1.54 veces en comparación con BF16, logrando un rendimiento casi sin pérdidas en diversas tareas, como el preentrenamiento y ajuste fino de Modelos de Lenguaje Grandes y el entrenamiento de Modelos de Lenguaje Visual. COAT también logra una aceleración del entrenamiento de extremo a extremo de 1.43 veces en comparación con BF16, desempeñándose al nivel o superando la aceleración de TransformerEngine. COAT permite el entrenamiento eficiente de modelos grandes con todos los parámetros en menos GPUs, y facilita duplicar el tamaño del lote en entornos de entrenamiento distribuido, proporcionando una solución práctica para escalar el entrenamiento de modelos a gran escala. El código está disponible en 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.