APOLLO: Memoria similar a SGD, rendimiento a nivel de 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

Resumen

Los modelos de lenguaje grandes (LLMs) son conocidos por ser intensivos en memoria durante el entrenamiento, especialmente con el popular optimizador AdamW. Esta carga de memoria hace necesario utilizar más o mejores GPUs o reducir los tamaños de lote, limitando la escalabilidad y el rendimiento del entrenamiento. Para abordar esto, se han propuesto varios optimizadores eficientes en memoria para reducir el uso de memoria del optimizador. Sin embargo, enfrentan desafíos críticos: (i) dependencia de operaciones costosas de SVD; (ii) compensaciones significativas de rendimiento en comparación con AdamW; y (iii) una sobrecarga de memoria del optimizador aún sustancial para mantener un rendimiento competitivo. En este trabajo, identificamos que la regla de adaptación de la tasa de aprendizaje de AdamW puede ser efectivamente simplificada como una actualización estructurada de la tasa de aprendizaje. Basándonos en esta idea, proponemos Escalado de Gradiente Aproximado para la Optimización Eficiente en Memoria de LLM (APOLLO), que aproxima el escalado de la tasa de aprendizaje utilizando un estado de optimizador de rango bajo auxiliar basado en proyección aleatoria pura. Esta regla de actualización estructurada de la tasa de aprendizaje hace que APOLLO sea altamente tolerante a reducciones adicionales de memoria al tiempo que ofrece un rendimiento de pre-entrenamiento comparable. Incluso su variante de rango-1, APOLLO-Mini, logra un rendimiento de pre-entrenamiento superior en comparación con AdamW con costos de memoria a nivel de SGD. Experimentos extensos demuestran que la serie APOLLO se desempeña igual o mejor que AdamW, al tiempo que logra mayores ahorros de memoria al eliminar casi por completo los estados de optimización de AdamW. Estos ahorros proporcionan beneficios significativos a nivel de sistema: (1) Mayor Rendimiento: 3 veces más rendimiento en una configuración de 8 GPUs A100-80GB en comparación con AdamW al admitir tamaños de lote 4 veces más grandes. (2) Mejora en la Escalabilidad del Modelo: Pre-entrenamiento de LLaMA-13B con DDP ingenuo en GPUs A100-80GB sin optimizaciones a nivel de sistema. (3) Pre-entrenamiento Amigable con GPUs de Bajo Rendimiento: Pre-entrenamiento de LLaMA-7B en una sola GPU utilizando menos de 12 GB de memoria con cuantificación 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.