COAT: Comprimere gli stati dell'ottimizzatore e l'attivazione per un addestramento efficiente in FP8.

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

October 25, 2024

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

cs.AI

Abstract

L'addestramento in FP8 è emerso come un metodo promettente per migliorare l'efficienza dell'addestramento. I framework esistenti accelerano l'addestramento applicando calcoli in FP8 ai livelli lineari lasciando gli stati dell'ottimizzatore e le attivazioni in una precisione più elevata, il che non ottimizza appieno l'utilizzo della memoria. Questo articolo introduce COAT (Compressing Optimizer States and Activations for FP8 Training), un nuovo framework di addestramento in FP8 progettato per ridurre significativamente l'impronta di memoria durante l'addestramento di modelli di grandi dimensioni. COAT affronta le limitazioni attuali attraverso due innovazioni chiave: (1) Espansione dell'Intervallo Dinamico, che allinea più strettamente le distribuzioni degli stati dell'ottimizzatore con l'intervallo di rappresentazione in FP8, riducendo così l'errore di quantizzazione, e (2) Quantizzazione delle Attivazioni a Granularità Mista, che ottimizza la memoria di attivazione utilizzando una combinazione di strategie di quantizzazione per tensore e per gruppo. Gli esperimenti dimostrano che COAT riduce efficacemente l'impronta di memoria di addestramento end-to-end di 1,54x rispetto a BF16, ottenendo prestazioni quasi senza perdite su varie attività, come il preaddestramento e il raffinamento di modelli di lingua di grandi dimensioni e l'addestramento di modelli di lingua per la visione. COAT raggiunge anche un aumento della velocità di addestramento end-to-end di 1,43x rispetto a BF16, eseguendo alla pari o superando l'aumento di velocità di TransformerEngine. COAT consente un addestramento efficiente di tutti i parametri di modelli di grandi dimensioni su meno GPU e facilita il raddoppio della dimensione del batch in ambienti di addestramento distribuito, offrendo una soluzione pratica per scalare l'addestramento di modelli su larga scala. Il codice è disponibile su 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.