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.
AdamW SVD operations learning rate adaptation rule Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO) auxiliary low-rank optimizer state pure random projection structured learning rate update pre-training performance APOLLO-Mini rank-1 variant LLaMA-13B naive DDP LLaMA-7B weight quantization