APOLLO: SGD-like Memory, AdamW-level Performance

APOLLO: SGD-like Memory, AdamW-level Performance

Speakers: Hanqing Zhu, PhD Student, UT Austin; Zhenyu Zhang; Wenyan Cong; Zhangyang Wang; Jinwon Lee

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of large language models (LLMs), the computational and memory demands for training these sophisticated architectures have become a significant bottleneck, limiting accessibility and innovation. Hanqing Zhu from UT Austin, alongside collaborators from UT Austin and AI, presented APOLLO, a groundbreaking solution designed to address this challenge at MLSys 2025. APOLLO introduces a novel memory-efficient training paradigm that promises the memory footprint of traditional Stochastic Gradient Descent (SGD) while consistently achieving or even surpassing the performance of AdamW, the de facto optimizer for modern LLMs.

The core problem APOLLO tackles is the exorbitant memory consumption of adaptive optimizers like Adam, primarily due to their need to store first and second-order momentum states for each model parameter. This overhead often doubles or triples the memory required beyond just the model weights and gradients, making it infeasible to train large models on readily available hardware. APOLLO's approach, rooted in rethinking the inherent redundancy within Adam's element-wise gradient scaling, leverages low-rank approximation and random projection to drastically cut down memory usage without compromising training stability or final model performance.

This talk is crucial for anyone involved in training or fine-tuning large machine learning models, from researchers and infrastructure teams to individual practitioners. By democratizing access to powerful LLM training, APOLLO has the potential to accelerate research, foster innovation in smaller labs, and enable more efficient utilization of existing GPU resources, fundamentally shifting the paradigm of how we approach large-scale model optimization.

Background

▶ Watch: APOLLO: SGD-like memory, Adam-level performance for LLMs (0:00)

The advent of large foundation models, particularly Large Language Models (LLMs), has revolutionized various fields, yet their training presents formidable challenges, especially concerning memory. A significant portion of this memory burden stems not from the model weights themselves, nor even their gradients, but from the optimizer states. For instance, training a Llama 7 billion parameter model requires approximately 58 gigabytes of memory in total. While model weights and gradients account for around 40 gigabytes, the Adam optimizer, the standard choice for stabilizing LLM training, demands an additional 28 gigabytes. This is because Adam, and its variants like AdamW, store both first-order momentum (mean of gradients) and second-order momentum (variance of squared gradients) for each parameter, effectively doubling the memory footprint dedicated solely to the optimizer states compared to the model size itself.

This excessive memory requirement has driven the development of various mitigation strategies. One common approach involves memory offloading, where portions of data (e.g., optimizer states) are moved from high-speed GPU memory (SRAM) to slower CPU memory (DRAM). While this reduces peak GPU memory usage, it introduces significant communication overhead, trading computation time for memory savings. Another avenue of research has explored algorithmic-level optimizations. These include low-rank based methods that approximate updates in a reduced dimensional space, quantization techniques such as 8-bit optimizers (e.g., bitsandbytes), or methods like Adam Mini that exploit optimizer redundancy through block-wise learning rate adoption.

However, these existing solutions often fall short. Many low-rank methods, particularly those relying on Singular Value Decomposition (SVD) or similar projections, incur substantial computational overhead. SVD operations scale cubically with the dimension (N^3), leading to dramatic throughput drops during training, sometimes even halting it entirely during initialization or periodic updates. For example, using SVD-based methods with Hugging Face libraries often triggers notifications about significant delays. Furthermore, while some methods avoid SVD, they often yield suboptimal performance compared to Adam, sometimes failing entirely for critical pre-training tasks (e.g., Aurora). Even when memory-efficient, they may still consume considerable memory; for instance, some require 1/4 of the original rank, and Adam Mini still necessitates storing full momentum.

The overarching research question that APOLLO seeks to answer is whether it's possible to develop a training solution that offers low overhead, achieves extreme memory efficiency akin to SGD, yet delivers the powerful performance and stability of Adam, which is indispensable for modern LLM training. The speaker posits that Adam, in its current form, is not optimally designed for LLMs and contains significant redundancy, motivating a deeper investigation into its internal mechanisms.

Key Findings

▶ Watch: Limitations of existing memory-efficient Adam solutions (1:50)

APOLLO represents a significant leap forward in memory-efficient deep learning, delivering a suite of compelling findings that challenge conventional wisdom about adaptive optimizers:

• SGD-like Memory, AdamW-level Performance: The most crucial finding is APOLLO's ability to achieve memory efficiency comparable to Stochastic Gradient Descent (SGD), which stores no optimizer states, while consistently matching or even exceeding the performance of Adam and AdamW across various training and fine-tuning tasks for large language models. This disproves the notion that significant memory reduction must come at the cost of performance.
• Record Memory Efficiency: APOLLO establishes a new benchmark for memory efficiency in LLM training. It can reduce the optimizer state memory footprint to negligible levels, approaching zero in its APOLLO Mini variant. When combined with weight quantization, APOLLO enables the training of a Llama 7 billion parameter model within a mere 20 gigabytes of total GPU memory. This unprecedented efficiency means large models can be trained on consumer-grade or older GPUs, such as NVIDIA Titan cards.
• Democratization of LLM Training: By drastically lowering the memory barrier, APOLLO makes large model training more accessible to a wider range of researchers, smaller teams, and individual practitioners who lack access to state-of-the-art, high-memory GPU clusters.
• Significant Throughput Gains: The memory savings translate directly into the ability to use larger batch sizes. This, in turn, leads to substantial throughput improvements, with APOLLO demonstrating up to a 3x faster training speed on the Llama 7 billion model compared to baselines constrained by memory.
• First Pre-training with Random Projection: APOLLO is the first method to successfully enable large-scale pre-training of LLMs using random projection for low-rank adaptive optimizers. This avoids the computational overhead of SVD-like methods while proving mathematically sound and empirically effective.
• Adam's Redundancy Confirmed: The research provides strong evidence that Adam's element-wise learning rate adaptation contains considerable redundancy. By structuring the gradient scanning factor (e.g., channel-wise, tensor-wise) and approximating it in a low-rank space, APOLLO demonstrates that this fine-grained control is not always necessary for optimal performance and can be significantly simplified.
• Robustness to Rank Reduction: APOLLO proves remarkably robust to aggressive rank reduction. For instance, it can effectively train a 60 million parameter model using a rank of just 4, and its Mini variant even uses rank 1, without detrimental impact on performance.

Technical Deep Dive

▶ Watch: APOLLO's core idea: Low-rank gradient scanning (5:55)

APOLLO's innovative approach stems from a fundamental re-evaluation of the Adam optimizer, specifically its mechanism for adaptive learning rate adjustment. The speaker begins by conceptualizing Adam as an SGD-like algorithm augmented with a gradient scanning factor. In essence, if this scanning factor were removed or set to identity, Adam would revert to pure SGD. The core of Adam's adaptivity lies in computing this element-wise scanning factor based on the historical momentum (first-order information) and variance (second-order information) of the gradients. This factor then scales the raw gradient to produce the update.

The critical question APOLLO poses is: Is this element-wise learning rate adoption truly necessary? The motivations for questioning this stem from several observations:

1. High-dimensional LLMs can be overly sensitive to noisy gradients. Highly precise, element-wise adaptive updates might trap the model in suboptimal local minima.
2. Prior work like Adam Mini has shown that block-wise learning rate adoption (e.g., within transformer blocks) can be effective, suggesting inherent similarities and redundancies in how different parts of a model's parameters should be adapted.

Initial explorations involved making the scanning factor structured, moving from element-wise to channel-wise or tensor-wise adaptation. Experiments on smaller models revealed that while early-stage training loss might diverge slightly, this could be easily managed with gradient clipping. Ultimately, the training loss converged similarly to full Adam, indicating that much of Adam's element-wise precision is redundant. However, this initial structuralization alone offered no memory benefits, as full momentum and variance still needed to be stored.

To achieve memory efficiency, APOLLO introduces its core innovation: approximating the gradient scanning in a low-rank space. The pipeline involves three key steps:

1. Gradient Projection: Instead of working with the full-rank gradient directly, APOLLO first projects the full gradient into a low-rank space. This compression step is crucial for memory reduction.
2. Low-Rank Adam: Within this reduced-dimensional space, APOLLO maintains low-rank momentum and low-rank variance. These are significantly smaller than their full-rank counterparts, leading to substantial memory savings. The gradient scanning factor is then computed using these low-rank states.
3. Reapplication to Full-Rank Gradient: Crucially, after obtaining the low-rank scanning factor, APOLLO reapplies it to the original full-rank gradient to compute the final parameter update. This distinguishes APOLLO from other low-rank methods like LoRA or QLoRA, which perform parameter updates entirely within the low-rank space. By updating the full-rank gradient, APOLLO ensures that the model's expressive capacity is fully utilized, explaining its superior performance compared to methods that restrict updates to a low-rank subspace.

In terms of memory benefits, this approach is transformative. For momentum and variance, instead of storing 2 M N parameters (where M is the model size and N is the dimension), APOLLO stores 2 M R parameters, where R is the chosen rank and R << N. This results in a drastic reduction in optimizer state memory.

A critical component of APOLLO's design is its choice of projection matrix. Unlike methods that rely on computationally expensive SVD, APOLLO uses a random projection. The speaker highlights that random projection offers several advantages:

• Zero Storage Overhead: The random projection matrix itself does not need to be stored; only a random seed is required, making it extremely memory efficient.
• Pre-training Compatibility: APOLLO is the first to demonstrate successful pre-training of LLMs with random projection for adaptive optimizers.
• Mathematical Elegance: The paper provides a theoretical proof showing that the scanning factor computed in the low-rank space and the original full-rank space have a consistent bound. This bound is simply a ratio of the original dimension and the chosen rank (N/R). This factor can be used to rescale the low-rank scanning factor, ensuring consistent behavior with full-rank Adam without introducing new hyperparameters. This "elegant factor" is empirically observed in training dynamics across various layers and stages.
• L2 Norm Preservation: Random projection naturally preserves the L2 norm of vectors with high probability, which contributes to the stability of the scaling factor.

APOLLO also demonstrates remarkable flexibility in rank reduction. It can achieve strong performance even with very low ranks, such as a rank of 4 for a 60 million parameter model. For even greater memory savings, APOLLO Mini is introduced, which operates with a rank of 1 and employs tensor-wise scanning to renormalize gradient scaling factors, further ensuring stable performance. The ability to use such aggressively low ranks without performance degradation underscores the inherent redundancy in Adam's full-rank adaptive mechanisms.

Experimental Setup & Results

▶ Watch: APOLLO's unique approach: Random projection, full-rank updates (6:45)

The evaluation of APOLLO spanned a comprehensive set of experiments, covering pre-training, fine-tuning, and system-level efficiency on various scales of large language models.

Models and Baselines:

The primary models used for evaluation were Llama models of different sizes, including the challenging Llama 7 billion parameter model. APOLLO's performance was rigorously compared against:

• Standard Adam/AdamW: The strong baseline, often considered the gold standard for LLM training. Learning rates for Adam were extensively swept to ensure a robust baseline comparison.
• 8-bit Adam versions: Such as those implemented in libraries like bitsandbytes, which are popular for their memory efficiency.
• Other memory-efficient optimizers: Including LoRA (low-rank adaptation) and Aurora, which were often found to be inferior, especially for pre-training.

Metrics:

Key performance indicators included:

• Memory Footprint: Specifically, the memory consumed by optimizer states and total GPU memory usage.
• Training Loss: To assess convergence and stability during pre-training.
• Fine-tuning Performance: Evaluated on a variety of downstream tasks (details in the full paper) to ensure generalization capabilities.
• Throughput (Samples/second): To quantify the system-level speed benefits.

Headline Numbers and Key Results:

1. Pre-training Scalability (Llama 7B):

• APOLLO, even with very small ranks (e.g., rank 4 for Llama 7B, or rank 1 for APOLLO Mini), demonstrated performance that was on par with or even better than Adam. This was a significant finding, as it dispelled the notion that cutting memory must compromise training quality.
• Crucially, APOLLO was the first solution capable of truly finishing the pre-training of a Llama 7 billion model within a fixed window (e.g., two weeks) without resorting to complex system-level techniques like model sharding. This was directly enabled by the substantial memory savings, which freed up resources for activations and larger batch sizes.

1. Fine-tuning Performance:

• Across numerous fine-tuning tasks, APOLLO consistently matched or outperformed Adam. It proved to be a highly memory-efficient, full-parameter fine-tuning solution, demonstrating its versatility beyond just pre-training.

1. Record Memory Efficiency:

• APOLLO achieved a new record in memory efficiency. It drastically reduced the optimizer part of the memory footprint, making it very small, and even effectively zero in the APOLLO Mini variant.
• When combined with weight quantization (e.g., 8-bit weights), APOLLO enabled the training of the Llama 7 billion model within a total of 20 gigabytes of GPU memory. This is a monumental achievement, allowing users with older or consumer-grade GPUs, such as a NVIDIA Titan, to train such large models from scratch.

1. System-Level Throughput Benefits:

• The memory savings directly translated into the ability to use larger batch sizes. This, in turn, resulted in significant throughput improvements. On the Llama 7 billion model, APOLLO achieved 3x better throughput compared to memory-constrained baselines.

Ablations and Observations:

• The performance of APOLLO was found to be not very sensitive to rank reduction. Even aggressively low ranks (e.g., rank 4 for a 60 million parameter model) did not significantly degrade performance, highlighting the robustness of the low-rank approximation.
• The speaker referenced earlier work on Adam Mini, which showed that a "50% drop" (presumably in some aspect of adaptivity or memory) did not hurt performance, providing further context and empirical support for APOLLO's premise regarding Adam's redundancy.
• Regarding the choice of random projection versus SVD, experiments showed that for pre-training, both performed very similarly. However, for fine-tuning, SVD-based projection could sometimes yield slightly better results, but at the cost of its well-known computational overhead. The speaker emphasized the elegance of random projection's bound and its non-hyperparameter nature as key advantages.

Overall, the experimental results unequivocally demonstrate APOLLO's capability to deliver state-of-the-art performance with unprecedented memory efficiency and significant throughput gains, making large language model training substantially more accessible.

Practical Implications

▶ Watch: Empirical validation of APOLLO's theoretical bound (8:40)

The introduction of APOLLO carries profound practical implications for the entire machine learning ecosystem, from individual researchers to large-scale infrastructure teams.

1. Democratization of LLM Training: Perhaps the most significant impact is the democratization of large language model training. By drastically reducing the memory requirements (e.g., training Llama 7B in 20GB), APOLLO enables individuals and smaller research groups to train or fine-tune powerful LLMs on readily available hardware, such as older or consumer-grade GPUs like NVIDIA Titan series cards. This lowers the barrier to entry, fostering innovation and broadening participation in LLM research and development beyond well-funded institutions.

2. Optimized Resource Utilization: For organizations with existing GPU clusters, APOLLO allows for more efficient utilization of hardware. Memory that would typically be consumed by optimizer states can now be repurposed. This means:

• Larger Batch Sizes: The primary benefit observed is the ability to use significantly larger batch sizes, leading to up to 3x higher training throughput on models like Llama 7B. Faster training cycles accelerate experimentation and model iteration.
• Larger Models: Teams can train even larger models than previously possible on their current hardware, pushing the boundaries of what's feasible without needing to acquire more expensive, high-memory GPUs.
• Longer Context Windows: The freed memory can also be allocated to store larger activations, enabling the use of longer sequence lengths or context windows during training, which is crucial for many advanced LLM applications.

3. Simplified Training Infrastructure: APOLLO's efficiency can simplify the complexity of training infrastructure. The ability to train a Llama 7B model within two weeks without complex model sharding or distributed training techniques (which often introduce their own overhead and engineering challenges) streamlines the development process. This reduces the operational burden on infrastructure teams and allows model builders to focus more on algorithmic advancements rather than system-level optimizations.

4. Reduced Carbon Footprint: While not explicitly stated, increased training efficiency and the ability to utilize older, less power-hungry hardware can indirectly contribute to a reduced carbon footprint associated with large-scale AI training, aligning with growing concerns about sustainable AI.

Tradeoffs and Limitations:

The talk acknowledges that while APOLLO's random projection is highly effective for pre-training and offers significant practical advantages, alternative projection methods like SVD might sometimes yield slightly better performance for specific fine-tuning tasks. However, the speaker emphasizes that the computational cost of SVD (scaling at N^3) makes it impractical for large-scale pre-training, where random projection's simplicity and theoretical guarantees (like L2 norm preservation and the elegant scaling bound) make it superior. The "elegant factor" for rescaling gradients, derived from the ratio of original dimension to rank, is a key practical benefit as it is not a hyperparameter that needs tuning. While APOLLO significantly reduces optimizer memory, it still requires model weights and gradients, meaning a base level of GPU memory will always be necessary.

In essence, APOLLO provides a powerful tool that makes high-performance LLM training accessible, faster, and more resource-efficient, fundamentally changing the economics and logistics of working with cutting-edge AI models.

Key Takeaways

• Memory-Efficient Training Paradigm: APOLLO enables SGD-like memory usage for optimizer states while achieving AdamW-level or superior performance in large language model training.
• Record-Breaking Memory Footprint: It allows training of Llama 7 billion parameter models within 20GB of total GPU memory (with weight quantization), making powerful LLM training accessible on consumer-grade or older GPUs like NVIDIA Titan.
• Low-Rank Adaptive Optimization: APOLLO leverages low-rank approximation of gradient scanning, computed efficiently via random projection, to drastically reduce the storage required for momentum and variance states.
• Significant Throughput Gains: Memory savings translate directly to the ability to use larger batch sizes, resulting in up to 3x faster training throughput on Llama 7B models.
• Democratization of LLMs: The reduced hardware requirements democratize access to large model training, empowering more researchers and practitioners globally.
• Seamless Integration: APOLLO is readily available and can be integrated into existing workflows via pip install or through Hugging Face.

About the Speaker(s)

Hanqing Zhu is a PhD student at UT Austin and was the primary presenter of the APOLLO work at MLSys 2025. This research is a joint effort between UT Austin and AI. The other listed speakers, Zhenyu Zhang, Wenyan Cong, Zhangyang Wang, and Jinwon Lee, contributed to this collaborative project.