Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MicroAdam: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence

Ionut-Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtic, Thomas Robert, Peter Richtarik, Dan Alistarh

TL;DR

It is shown that MicroAdam can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical convergence competitive to that of the uncompressed Adam baseline, with lower memory usage and similar running time.

Abstract

We propose a new variant of the Adam optimizer called MicroAdam that specifically minimizes memory overheads, while maintaining theoretical convergence guarantees. We achieve this by compressing the gradient information before it is fed into the optimizer state, thereby reducing its memory footprint significantly. We control the resulting compression error via a novel instance of the classical \emph{error feedback} mechanism from distributed optimization in which *the error correction information is itself compressed* to allow for practical memory gains. We prove that the resulting approach maintains theoretical convergence guarantees competitive to those of AMSGrad, while providing good practical performance. Specifically, we show that MicroAdam can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical convergence competitive to that of the uncompressed Adam baseline, with lower memory usage and similar running time. Our code is available at https://github.com/IST-DASLab/MicroAdam.

MicroAdam: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence

TL;DR

It is shown that MicroAdam can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical convergence competitive to that of the uncompressed Adam baseline, with lower memory usage and similar running time.

Abstract

We propose a new variant of the Adam optimizer called MicroAdam that specifically minimizes memory overheads, while maintaining theoretical convergence guarantees. We achieve this by compressing the gradient information before it is fed into the optimizer state, thereby reducing its memory footprint significantly. We control the resulting compression error via a novel instance of the classical \emph{error feedback} mechanism from distributed optimization in which *the error correction information is itself compressed* to allow for practical memory gains. We prove that the resulting approach maintains theoretical convergence guarantees competitive to those of AMSGrad, while providing good practical performance. Specifically, we show that MicroAdam can be implemented efficiently on GPUs: on both million-scale (BERT) and billion-scale (LLaMA) models, MicroAdam provides practical convergence competitive to that of the uncompressed Adam baseline, with lower memory usage and similar running time. Our code is available at https://github.com/IST-DASLab/MicroAdam.
Paper Structure (56 sections, 19 theorems, 98 equations, 9 figures, 4 tables, 4 algorithms)

This paper contains 56 sections, 19 theorems, 98 equations, 9 figures, 4 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Related Work
  4. The MicroAdam Algorithm
  5. Notation.
  6. Algorithm Description.
  7. Properties and Limitations.
  8. Dynamic Statistics.
  9. Algorithm Intuition.
  10. Efficient Implementation
  11. Accumulator $a_t$.
  12. Top-K.
  13. Quantization metadata.
  14. Practical memory usage.
  15. Memory footprint analysis for the optimizer states and comparison with other methods
  16. ...and 41 more sections

Key Result

Lemma 1

Consider Algorithm algorithm:procedures with randomized rounding, i.e., for a vector $x\in\mathbb{R}^d$ with $\delta = \min_i x_i$ and $\Delta = \max_i x_i$, let $\hat{x}_i \mathrel{\mathop:}= \lfloor \frac{x_i-\delta}{u} + \xi \rfloor u + \delta$ be the $i$-th coordinate of the quantized vector $\h

Figures (9)

  • Figure 1: Optimization trajectories of Adam, TopK-Adam and TopK-Adam with EF applied on the Rosenbrock function $f(x,y)=(1-x)^2+100(y-x^2)^2$ starting from $(x_0,y_0)=(-\frac{1}{2},1)$. Notice the extremely "jagged" profile of TopK-Adam without EF, and the recovered convergence when EF is added.
  • Figure 2: Training curves for BERT-Base on GLUE/MNLI
  • Figure 3: Training curves for BERT-Large on GLUE/MNLI
  • Figure 4: Training curves for OPT-1.3B on GLUE/MNLI
  • Figure 5: Training curves for Llama-2 7B on GSM-8k
  • ...and 4 more figures

Theorems & Definitions (33)

  • Lemma 1
  • Theorem 1
  • Theorem 2
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • proof
  • Lemma 4
  • ...and 23 more