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