COAT: Compressão de Estados do Otimizador e Ativação para Treinamento Eficiente em Memória 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

Resumo

O treinamento em FP8 surgiu como um método promissor para melhorar a eficiência do treinamento. Estruturas existentes aceleram o treinamento aplicando cálculos em FP8 às camadas lineares, enquanto deixam os estados do otimizador e ativações em maior precisão, o que falha em otimizar totalmente o uso da memória. Este artigo apresenta o COAT (Compressão de Estados do Otimizador e Ativações para Treinamento em FP8), um novo framework de treinamento em FP8 projetado para reduzir significativamente a pegada de memória ao treinar modelos grandes. O COAT aborda as limitações atuais por meio de duas inovações-chave: (1) Expansão de Faixa Dinâmica, que alinha mais de perto as distribuições de estados do otimizador com a faixa de representação do FP8, reduzindo assim o erro de quantização, e (2) Quantização de Ativação de Granularidade Mista, que otimiza a memória de ativação usando uma combinação de estratégias de quantização por tensor e por grupo. Experimentos demonstram que o COAT reduz efetivamente a pegada de memória de treinamento de ponta a ponta em 1,54x em comparação com o BF16, enquanto alcança um desempenho quase sem perdas em várias tarefas, como pré-treinamento de Modelos de Linguagem Grandes e ajuste fino, e treinamento de Modelos de Linguagem Visual. O COAT também alcança um aumento de velocidade de treinamento de ponta a ponta de 1,43x em comparação com o BF16, atuando em pé de igualdade ou superando o aumento de velocidade do TransformerEngine. O COAT possibilita o treinamento eficiente de todos os parâmetros de modelos grandes em menos GPUs e facilita a duplicação do tamanho do lote em configurações de treinamento distribuído, fornecendo uma solução prática para a escalabilidade do treinamento de modelos em larga escala. O código está disponível em 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.