APOLLO: SGD-like Memory, AdamW-level Performance
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.
675667e201ef91bb7046f9f8
Hanqing Zhu
admin_assigned
2024-12-09T15:16:08.243Z
64adab87692be9d12f81957d
/avatars/dbb738c8e78769513d8d699e62d6e49b.svg
Hanqing Zhu
user
hanqing666
675667e201ef91bb7046f9f9
Zhenyu Zhang
claimed_verified
2024-12-09T09:46:20.874Z
649c888f67fd6c6aa97e5f85
/avatars/9967b729916d1128773102797fed1673.svg
Zhenyu Zhang
user
Kyriection
675667e201ef91bb7046f9fa
Wenyan Cong
admin_assigned
2024-12-09T15:16:16.157Z
63884e0a3143f4706312f8eb
/avatars/cd9262a8b5b6e273c629b4d1c76a6d46.svg
Wenyan Cong
user
FOGmiaow
675667e201ef91bb7046f9fb
Xi Liu
admin_assigned
2024-12-09T15:16:22.759Z
665a92e30b1b9f88e185d74a
/avatars/1a606cec39963b5f11f46fe3952d93c6.svg
Xi Liu
user
XiLiu814
675667e201ef91bb7046f9fc
Sem Park
675667e201ef91bb7046f9fd
Vikas Chandra
admin_assigned
2024-12-09T15:16:46.417Z
6566681f7b5ed0735812af32
/avatars/424905abb2973954b0850c592743b6fb.svg
Vikas Chandra
user
vchandra
675667e201ef91bb7046f9fe
Bo Long
675667e201ef91bb7046f9ff
David Z. Pan
675667e201ef91bb7046fa00
Zhangyang Wang
675667e201ef91bb7046fa01
Jinwon Lee
APOLLO: SGD-like Memory, AdamW-level Performance
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.
APOLLO, an approximated gradient scaling optimizer for LLMs, reduces memory usage without compromising performance, enhancing throughput and scalability.
AdamW memory-efficient optimizers SVD operations structured learning rate update random projection APOLLO APOLLO-Mini throughput model scalability weight quantization