APOLLO: Memória semelhante ao SGD, Desempenho ao nível do AdamW

APOLLO: SGD-like Memory, AdamW-level Performance

December 6, 2024

Autores: Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z. Pan, Zhangyang Wang, Jinwon Lee

cs.AI

Resumo

Os grandes modelos de linguagem (LLMs) são conhecidos por serem intensivos em memória durante o treinamento, especialmente com o popular otimizador AdamW. Esse ônus de memória requer o uso de mais ou de GPUs mais avançadas, ou a redução dos tamanhos dos lotes, limitando a escalabilidade e a taxa de processamento do treinamento. Para lidar com isso, vários otimizadores eficientes em memória foram propostos para reduzir o uso de memória do otimizador. No entanto, eles enfrentam desafios críticos: (i) dependência de operações SVD custosas; (ii) significativos compromissos de desempenho em comparação com o AdamW; e (iii) ainda uma sobrecarga substancial de memória do otimizador para manter um desempenho competitivo. Neste trabalho, identificamos que a regra de adaptação da taxa de aprendizado do AdamW pode ser efetivamente simplificada como uma atualização estruturada da taxa de aprendizado. Com base nessa percepção, propomos Escalonamento de Gradiente Aproximado para Otimização Eficiente em Memória de LLM (APOLLO), que aproxima o escalonamento da taxa de aprendizado usando um estado auxiliar de otimizador de baixa classificação baseado em projeção aleatória pura. Essa regra de atualização estruturada da taxa de aprendizado torna o APOLLO altamente tolerante a reduções adicionais de memória, ao mesmo tempo em que oferece um desempenho de pré-treinamento comparável. Mesmo sua variante de classificação 1, APOLLO-Mini, alcança um desempenho de pré-treinamento superior em comparação com o AdamW com custos de memória no nível do SGD. Experimentos extensivos demonstram que a série APOLLO tem desempenho equivalente ou melhor que o AdamW, enquanto alcança maiores economias de memória ao quase eliminar os estados de otimização do AdamW. Essas economias proporcionam benefícios significativos em nível de sistema: (1) Aumento de Taxa de Processamento: 3x a taxa de processamento em uma configuração de 8 GPUs A100-80GB em comparação com o AdamW, suportando tamanhos de lote 4x maiores. (2) Melhoria na Escalabilidade do Modelo: Pré-treinamento do LLaMA-13B com DDP ingênuo em GPUs A100-80GB sem otimizações em nível de sistema. (3) Pré-treinamento Amigável para GPUs de Baixo Desempenho: Pré-treinamento do LLaMA-7B em uma única GPU usando menos de 12 GB de memória com quantização de pesos.

English

Large language models (LLMs) are notoriously memory-intensive during training, particularly with the popular AdamW optimizer. This memory burden necessitates using more or higher-end GPUs or reducing batch sizes, limiting training scalability and throughput. To address this, various memory-efficient optimizers have been proposed to reduce optimizer memory usage. However, they face critical challenges: (i) reliance on costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial optimizer memory overhead to maintain competitive performance. In this work, we identify that AdamW's learning rate adaptation rule can be effectively coarsened as a structured learning rate update. Based on this insight, we propose Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO), which approximates learning rate scaling using an auxiliary low-rank optimizer state based on pure random projection. This structured learning rate update rule makes APOLLO highly tolerant to further memory reductions while delivering comparable pre-training performance. Even its rank-1 variant, APOLLO-Mini, achieves superior pre-training performance compared to AdamW with SGD-level memory costs. Extensive experiments demonstrate that the APOLLO series performs on-par with or better than AdamW, while achieving greater memory savings by nearly eliminating the optimization states of AdamW. These savings provide significant system-level benefits: (1) Enhanced Throughput: 3x throughput on an 8xA100-80GB setup compared to AdamW by supporting 4x larger batch sizes. (2) Improved Model Scalability: Pre-training LLaMA-13B with naive DDP on A100-80GB GPUs without system-level optimizations. (3) Low-End GPU Friendly Pre-training: Pre-training LLaMA-7B on a single GPU using less than 12 GB of memory with weight quantization.