APOLLO: Memoria simile a SGD, prestazioni di livello AdamW

APOLLO: SGD-like Memory, AdamW-level Performance

December 6, 2024

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

cs.AI

Abstract

I grandi modelli linguistici (LLM) sono notoriamente intensivi in termini di memoria durante l'addestramento, specialmente con l'ottimizzatore AdamW popolare. Questo onere sulla memoria richiede l'uso di più o di GPU di fascia alta o la riduzione delle dimensioni dei batch, limitando la scalabilità e la capacità di addestramento. Per affrontare ciò, sono state proposte varie ottimizzazioni efficienti in termini di memoria per ridurre l'uso della memoria dell'ottimizzatore. Tuttavia, affrontano sfide critiche: (i) dipendenza da costose operazioni di SVD; (ii) significativi compromessi sulle prestazioni rispetto ad AdamW; e (iii) ancora un notevole sovraccarico di memoria dell'ottimizzatore per mantenere prestazioni competitive. In questo lavoro, identifichiamo che la regola di adattamento del tasso di apprendimento di AdamW può essere efficacemente approssimata come un aggiornamento strutturato del tasso di apprendimento. Basandoci su questa intuizione, proponiamo la Scalatura del Gradiente Approssimata per l'Ottimizzazione Efficientemente Memoria-Intensiva dei LLM (APOLLO), che approssima la scalatura del tasso di apprendimento utilizzando uno stato dell'ottimizzatore ausiliario a basso rango basato su una pura proiezione casuale. Questa regola di aggiornamento strutturato del tasso di apprendimento rende APOLLO altamente tollerante a ulteriori riduzioni di memoria pur offrendo prestazioni di pre-addestramento comparabili. Anche la sua variante di rango-1, APOLLO-Mini, raggiunge prestazioni di pre-addestramento superiori rispetto ad AdamW con costi di memoria a livello di SGD. Esperimenti estesi dimostrano che la serie APOLLO si comporta alla pari o meglio di AdamW, ottenendo nel contempo maggiori risparmi di memoria eliminando quasi completamente gli stati di ottimizzazione di AdamW. Questi risparmi forniscono significativi vantaggi a livello di sistema: (1) Maggiore Capacità: 3 volte la capacità su una configurazione 8xA100-80GB rispetto ad AdamW supportando dimensioni di batch 4 volte più grandi. (2) Miglior Scalabilità del Modello: Pre-addestramento di LLaMA-13B con DDP ingenuo su GPU A100-80GB senza ottimizzazioni a livello di sistema. (3) Pre-addestramento Amichevole per GPU di Fascia Bassa: Pre-addestramento di LLaMA-7B su una singola GPU utilizzando meno di 12 GB di memoria con quantizzazione dei pesi.

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.