COAT: Сжатие состояний оптимизатора и активации для обучения с плавающей запятой 8-битной точности с учетом эффективности памяти

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

October 25, 2024

Авторы: Haocheng Xi, Han Cai, Ligeng Zhu, Yao Lu, Kurt Keutzer, Jianfei Chen, Song Han

cs.AI

Аннотация

Обучение с использованием FP8 стало перспективным методом для улучшения эффективности обучения. Существующие фреймворки ускоряют обучение, применяя вычисления FP8 к линейным слоям, оставляя при этом состояния оптимизатора и активации в более высокой точности, что не позволяет полностью оптимизировать использование памяти. В данной статье представлен COAT (Compressing Optimizer States and Activations for FP8 Training), новый фреймворк обучения с использованием FP8, разработанный для значительного снижения объема памяти при обучении больших моделей. COAT решает текущие ограничения благодаря двум ключевым инновациям: (1) Расширение динамического диапазона, которое выравнивает распределение состояний оптимизатора ближе к диапазону представления FP8, тем самым уменьшая ошибку квантования, и (2) Квантование активации смешанной гранулярности, которое оптимизирует использование памяти активации с помощью комбинации стратегий квантования на уровне тензора и на уровне группы. Эксперименты показывают, что COAT эффективно снижает общий объем памяти для обучения на 1,54 раза по сравнению с BF16, сохраняя практически без потерь производительность на различных задачах, таких как предварительное обучение и настройка больших языковых моделей и обучение моделей видео-языка. COAT также достигает ускорения общего времени обучения на 1,43 раза по сравнению с BF16, демонстрируя производительность на уровне или превосходящую ускорение TransformerEngine. COAT позволяет эффективное обучение полных параметров больших моделей на меньшем количестве GPU и облегчает удвоение размера пакета в распределенных настройках обучения, предоставляя практическое решение для масштабирования обучения моделей большого масштаба. Код доступен по ссылке 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.