APOLLO: SGD-ähnlicher Speicher, Leistung auf AdamW-Niveau

APOLLO: SGD-like Memory, AdamW-level Performance

December 6, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) sind bekanntermaßen speicherintensiv während des Trainings, insbesondere mit dem beliebten AdamW-Optimierer. Diese Speicherbelastung erfordert den Einsatz von mehr oder leistungsstärkeren GPUs oder die Reduzierung von Batch-Größen, was die Trainierbarkeit und Durchsatz begrenzt. Um diesem Problem zu begegnen, wurden verschiedene speicher-effiziente Optimierer vorgeschlagen, um den Speicherverbrauch des Optimierers zu reduzieren. Allerdings stehen sie vor kritischen Herausforderungen: (i) Abhängigkeit von kostspieligen SVD-Operationen; (ii) signifikante Leistungskompromisse im Vergleich zu AdamW; und (iii) immer noch erheblicher Speicherüberkopf des Optimierers, um eine wettbewerbsfähige Leistung aufrechtzuerhalten. In dieser Arbeit identifizieren wir, dass die Lernratenanpassungsregel von AdamW effektiv als strukturiertes Lernratenupdate vereinfacht werden kann. Basierend auf dieser Erkenntnis schlagen wir die Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO) vor, die die Skalierung der Lernrate unter Verwendung eines Hilfs-Optimizerzustands mit niedriger Rangordnung basierend auf rein zufälliger Projektion approximiert. Diese Regel für strukturiertes Lernratenupdate macht APOLLO äußerst tolerant gegenüber weiteren Speicherreduktionen, während sie eine vergleichbare Vor-Trainingsleistung liefert. Selbst die Rang-1-Variante, APOLLO-Mini, erzielt eine überlegene Vor-Trainingsleistung im Vergleich zu AdamW mit Speicherkosten auf SGD-Niveau. Umfangreiche Experimente zeigen, dass die APOLLO-Serie auf Augenhöhe oder besser als AdamW abschneidet, während sie durch nahezu vollständige Beseitigung der Optimierungszustände von AdamW größere Speichereinsparungen erzielt. Diese Einsparungen bieten signifikante systemweite Vorteile: (1) Verbesserte Durchsatzleistung: 3-facher Durchsatz auf einem 8xA100-80GB-Setup im Vergleich zu AdamW durch Unterstützung von 4-fach größeren Batch-Größen. (2) Verbesserte Modell-Skalierbarkeit: Vor-Trainieren von LLaMA-13B mit naivem DDP auf A100-80GB-GPUs ohne systemweite Optimierungen. (3) Vor-Trainieren für Low-End-GPUs: Vor-Trainieren von LLaMA-7B auf einer einzelnen GPU unter Verwendung von weniger als 12 GB Speicher mit Gewichtsquantisierung.

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.