APOLLO : Mémoire similaire à SGD, performances de niveau AdamW

APOLLO: SGD-like Memory, AdamW-level Performance

December 6, 2024

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

cs.AI

Résumé

Les grands modèles de langage (LLM) sont connus pour leur consommation de mémoire élevée pendant l'entraînement, en particulier avec l'optimiseur AdamW populaire. Ce fardeau de mémoire nécessite l'utilisation de plus de GPU ou de GPU haut de gamme, ou la réduction des tailles de lot, limitant ainsi la scalabilité et le débit de l'entraînement. Pour remédier à cela, divers optimiseurs économes en mémoire ont été proposés pour réduire l'utilisation de mémoire de l'optimiseur. Cependant, ils rencontrent des défis critiques : (i) dépendance à des opérations SVD coûteuses ; (ii) compromis significatifs en termes de performances par rapport à AdamW ; et (iii) surcharge de mémoire d'optimiseur encore substantielle pour maintenir des performances compétitives. Dans ce travail, nous identifions que la règle d'adaptation du taux d'apprentissage d'AdamW peut être efficacement simplifiée en tant que mise à jour structurée du taux d'apprentissage. Sur la base de cette observation, nous proposons l'Échelle de Gradient Approximée pour l'Optimisation Économe en Mémoire des LLM (APOLLO), qui approxime la mise à l'échelle du taux d'apprentissage en utilisant un état auxiliaire d'optimiseur de rang faible basé sur une projection aléatoire pure. Cette règle de mise à jour structurée du taux d'apprentissage rend APOLLO très tolérant aux réductions de mémoire supplémentaires tout en offrant des performances de pré-entraînement comparables. Même sa variante de rang 1, APOLLO-Mini, atteint des performances de pré-entraînement supérieures par rapport à AdamW avec des coûts mémoire de niveau SGD. Des expériences approfondies démontrent que la série APOLLO se comporte au même niveau, voire mieux qu'AdamW, tout en réalisant des économies de mémoire plus importantes en éliminant presque entièrement les états d'optimisation d'AdamW. Ces économies offrent des avantages significatifs au niveau du système : (1) Débit Amélioré : un débit 3 fois supérieur sur une configuration 8xA100-80GB par rapport à AdamW en prenant en charge des tailles de lot 4 fois plus grandes. (2) Scalabilité du Modèle Améliorée : Pré-entraînement LLaMA-13B avec DDP naïf sur des GPU A100-80GB sans optimisations au niveau du système. (3) Pré-entraînement Amical aux GPU Bas de Gamme : Pré-entraînement LLaMA-7B sur un seul GPU en utilisant moins de 12 Go de mémoire avec quantification des poids.

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.