APOLLO: SGD-achtig Geheugen, Prestaties op het niveau van 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

Samenvatting

Grote taalmodellen (LLM's) staan bekend om hun intensieve geheugengebruik tijdens training, vooral met de populaire AdamW-optimalisator. Deze geheugenlast vereist het gebruik van meer of geavanceerdere GPU's of het verkleinen van batchgroottes, wat de schaalbaarheid en doorvoer van training beperkt. Om dit aan te pakken, zijn verschillende geheugenefficiënte optimalisatoren voorgesteld om het geheugengebruik van de optimalisator te verminderen. Ze worden echter geconfronteerd met kritieke uitdagingen: (i) afhankelijkheid van kostbare SVD-operaties; (ii) aanzienlijke prestatieafwegingen in vergelijking met AdamW; en (iii) nog steeds aanzienlijke geheugenoverhead van de optimalisator om concurrerende prestaties te behouden. In dit werk identificeren we dat de leersnelheidsaanpassingsregel van AdamW effectief kan worden verfijnd als een gestructureerde leersnelheidsupdate. Op basis van deze inzichten stellen we Approximated Gradient Scaling voor voor Geheugenefficiënte LLM-optimalisatie (APOLLO), die leersnelheidsverandering benadert met behulp van een hulpoptimalisatorstaat met lage rang gebaseerd op pure willekeurige projectie. Deze gestructureerde leersnelheidsupdate-regel maakt APOLLO zeer tolerant voor verdere geheugenvermindering terwijl het vergelijkbare pre-training prestaties levert. Zelfs de rank-1 variant, APOLLO-Mini, behaalt superieure pre-training prestaties in vergelijking met AdamW met SGD-niveau geheugenkosten. Uitgebreide experimenten tonen aan dat de APOLLO-serie op gelijke hoogte presteert of beter dan AdamW, terwijl het grotere geheugenbesparingen behaalt door bijna de optimalisatorstaten van AdamW te elimineren. Deze besparingen bieden aanzienlijke systeemvoordelen: (1) Verbeterde Doorvoer: 3x doorvoer op een 8xA100-80GB opstelling in vergelijking met AdamW door ondersteuning van 4x grotere batchgroottes. (2) Verbeterde Model Schaalbaarheid: Pre-training LLaMA-13B met naïeve DDP op A100-80GB GPU's zonder systeemniveau-optimalisaties. (3) Pre-training Vriendelijk voor GPU's met lage specificaties: Pre-training LLaMA-7B op een enkele GPU met minder dan 12 GB geheugen met gewichtskwantisering.

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.