APOLLO: Память, аналогичная SGD, производительность на уровне AdamW

APOLLO: SGD-like Memory, AdamW-level Performance

December 6, 2024

Авторы: Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z. Pan, Zhangyang Wang, Jinwon Lee

cs.AI

Аннотация

Большие языковые модели (LLM) известны своим высоким потреблением памяти во время обучения, особенно при использовании популярного оптимизатора AdamW. Эта нагрузка на память требует использования большего количества или более мощных графических процессоров или уменьшения размеров пакетов, что ограничивает масштабируемость и производительность обучения. Для решения этой проблемы были предложены различные оптимизаторы, эффективные с точки зрения использования памяти оптимизаторов, однако они сталкиваются с критическими вызовами: (i) зависимость от дорогостоящих операций SVD; (ii) значительные компромиссы в производительности по сравнению с AdamW; и (iii) все еще значительные накладные расходы на память оптимизатора для поддержания конкурентоспособной производительности. В данной работе мы выявляем, что правило адаптации скорости обучения AdamW может быть эффективно упрощено как структурированное обновление скорости обучения. Исходя из этого наблюдения, мы предлагаем метод Approximated Gradient Scaling для оптимизации памяти LLM (APOLLO), который аппроксимирует масштабирование скорости обучения с использованием вспомогательного оптимизатора низкого ранга на основе чистой случайной проекции. Это структурированное правило обновления скорости обучения делает APOLLO очень устойчивым к дальнейшему снижению потребления памяти, обеспечивая сравнимую производительность предварительного обучения. Даже его вариант с рангом 1, APOLLO-Mini, достигает более высокой производительности предварительного обучения по сравнению с AdamW с расходами памяти на уровне SGD. Обширные эксперименты показывают, что серия APOLLO выполняет задачу на уровне или лучше, чем AdamW, обеспечивая при этом большие экономии памяти путем почти полного устранения оптимизационных состояний AdamW. Эти экономии обеспечивают значительные преимущества на уровне системы: (1) Увеличенная производительность: в 3 раза большая производительность на установке 8xA100-80GB по сравнению с AdamW за счет поддержки пакетов в 4 раза большего размера. (2) Улучшенная масштабируемость модели: Предварительное обучение LLaMA-13B с примитивным DDP на A100-80GB GPU без оптимизаций на уровне системы. (3) Предварительное обучение, дружественное к графическим процессорам низкого уровня: Предварительное обучение LLaMA-7B на одном GPU с использованием менее 12 ГБ памяти с квантованием весов.

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.