COAT: Comprimeren van Optimizer-states en Activatie voor Geheugenefficiënte FP8 Training

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

Samenvatting

FP8-training is naar voren gekomen als een veelbelovende methode om de trainings-efficiëntie te verbeteren. Bestaande frameworks versnellen de training door FP8-berekeningen toe te passen op lineaire lagen, terwijl ze de optimizer states en activaties in hogere precisie laten, wat niet volledig optimaal is voor geheugengebruik. Dit artikel introduceert COAT (Compressing Optimizer States and Activations for FP8 Training), een nieuw FP8-trainingsframework dat is ontworpen om aanzienlijk de geheugenfootprint te verkleinen bij het trainen van grote modellen. COAT pakt huidige beperkingen aan door middel van twee belangrijke innovaties: (1) Dynamische Bereikuitbreiding, die optimizer state distributies nauwer afstemt op het FP8-representatiebereik, waardoor kwantisatiefouten worden verminderd, en (2) Gemengde-Granulariteit Activatiekwantisatie, die activatiegeheugen optimaliseert door een combinatie van per-tensor en per-groep kwantisatiestrategieën. Experimenten tonen aan dat COAT effectief de end-to-end trainingsgeheugenfootprint met 1.54x verkleint in vergelijking met BF16, terwijl het bijna verliesloze prestaties behaalt bij verschillende taken, zoals vooraftraining en fijnafstemming van grote taalmodellen en training van visueel taalmodellen. COAT behaalt ook een 1.43x versnelling van de end-to-end training in vergelijking met BF16, waarbij het op gelijke hoogte presteert of TransformerEngine's versnelling overtreft. COAT maakt efficiënte training van volledige parameters van grote modellen mogelijk op minder GPU's en vergemakkelijkt het verdubbelen van de batchgrootte in gedistribueerde trainingsomgevingen, wat een praktische oplossing biedt voor het schalen van grootschalige modeltraining. De code is beschikbaar op 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.