FlashOptim: Optimizers for Memory Efficient Training

AI summary

Overview

• Research area: Memory-efficient neural network training, specifically low-precision optimizers and optimizer state compression.
• Technical level: Intermediate to Advanced. Readers need familiarity with mixed-precision training, floating-point representation (ULP, BF16, FP16, FP32), and standard optimizer update rules (SGD, AdamW, Lion).
• Scope: The paper introduces FlashOptim, a set of optimizations that reduce per-parameter accelerator memory during training while preserving optimizer semantics and model quality, and validates them on image classification, LLM pretraining, and LLM finetuning.

What This Paper Is About

Standard mixed-precision training requires many bytes of accelerator memory per model parameter: the parameter itself, its gradient, and one or more optimizer state variables, each typically stored in 4 bytes. The authors state that training even a 7-billion-parameter model can require at least 112 GiB of accelerator memory under Adam, which is impractical for researchers with less than 100 GiB. FlashOptim targets this problem by compressing master weights and optimizer states with drop-in optimizer replacements that the authors report run just as fast as standard optimizers and produce no measurable quality degradation.

Key Contributions

1. Improved float splitting. Instead of materializing both a 32-bit master weight and a 16-bit downcast weight, the authors split each master weight into a low-precision weight plus an error-correction term stored in the optimizer. Their scheme supports either 8-bit or 16-bit error correction and achieves much lower reconstruction error for a given number of correction bits, allowing 24-bit master weights with no loss of model quality. The key insight is that the exponent bits of the correction term are wasted because the exponent of the error can be inferred from the downcast weight under round-to-nearest; they instead rescale the error by the unit in the last place (ULP) and store it as a b-bit integer.
2. Companded optimizer state quantization. Prior work showed optimizer states can be compressed to 8 bits per element with sufficient software complexity. The authors show this can be done with a one-line nonlinear preprocessing (companding) function before standard group-wise linear quantization, with separate functions designed for momentum and variance tensors.
3. Fused optimized kernels. FlashOptim is implemented as optimizer step kernels that fuse all compression and quantization operations into a single Triton kernel, preserving throughput during training.
4. Open-source release. The implementation is available as a PyTorch library at https://github.com/databricks/flashoptim.

Main Findings

• Memory reduction: FlashOptim reduces Adam from 16 to 7 bytes per parameter and SGD from 12 to 6 bytes, a reduction of over 50%. With gradient release, this drops further by 2 bytes, to 5 bytes for Adam and 4 bytes for SGD.
• Llama-3.1-8B finetuning: Peak memory drops from 175 GiB to 113 GiB (Figure 1; Table 4 reports 175.2 GiB to 112.9 GiB, a 36% reduction). Master weight memory falls from 29.9 GiB to 15.0 GiB (-50%) and optimizer state memory from 59.8 GiB to 23.4 GiB (-61%), while optimizer step time changes from 12.5 ms to 11.5 ms.
• Component ablations: Weight splitting alone halves master weight memory but adds 12% extra optimizer state (67.3 GiB vs. 59.8 GiB). Optimizer state quantization alone reduces optimizer state by approximately 73% (to 15.9 GiB), slightly less than 75% because of the FP16 scale factor overhead per group of 32 elements.
• Image classification (ResNet-50, ImageNet Top-1): Reference SGD 77.01 ± 0.02, FlashOptim 77.16 ± 0.09. Reference AdamW 75.51 ± 0.09, FlashOptim 75.67 ± 0.04.
• LLM finetuning (GSM8k, Llama-3.1-8B): Reference AdamW 75.09 ± 0.40, FlashOptim 74.98 ± 0.77.
• LLM pretraining (GPT-2 124M): Validation loss 3.263 ± 0.001 for AdamW vs. 3.265 ± 0.001 for FlashAdamW; 3.240 ± 0.002 for Lion vs. 3.240 ± 0.001 for FlashLion. Mean ICL accuracy 44.0 ± 0.4 for AdamW vs. 45.0 ± 1.0 for FlashAdamW; 44.0 ± 0.5 for Lion vs. 44.7 ± 0.5 for FlashLion. The paper reports FlashOptim scores within variance of reference optimizers across all metrics.
• Weight reconstruction: With BF16 weights and 16-bit ULP correction, the approach achieves near-zero error (<10⁻⁹) and perfect bitwise reconstruction in 99.92% of values. Storing the error in BF16 instead produces error >10⁻⁶, comparable to the 24-bit format. For FP16, the 32-bit ULP format perfectly reconstructs values in the normal range, and the 24-bit format improves worst-case error from 10⁻⁴ to under 10⁻⁶.
• Quantization error and stability: Companding reduces normalized MSE for momentum buffers and gives particularly large improvements for variance buffers across optimizers and datasets. In LLM pretraining, linear quantization of optimizer states causes rapid divergence, whereas companding maintains stable convergence.
• Checkpoint size: Standard Adam checkpoints require 12 bytes per parameter; FlashAdamW requires 5 bytes (2 for weights, 1 for correction, 1 for momentum, 1 for variance). For a 7B model, checkpoint size reduces from 84 GiB to 35 GiB.
• Distributed compatibility: Only the 16-bit parameters are all-gathered during forward and backward passes under FSDP; the correction term remains local with the optimizer states.

Methodology in Plain English

The authors modify the optimizer step by adding a prologue and an epilogue while leaving the update rule itself untouched. In the prologue, stored optimizer states are dequantized and the full-precision master weight is reconstructed from a low-precision weight plus an error-correction term. The standard optimizer update runs as usual. In the epilogue, the new optimizer state is quantized and the new master weight is split back into a low-precision weight and correction bits.

For weight storage, rather than keeping a 32-bit master weight and a redundant 16-bit downcast weight, they keep a 16-bit weight and a small integer correction. Because round-to-nearest guarantees the original value lies within half a ULP of the stored value, the exponent of the error is already known; the error only needs to be rescaled into that interval and stored as an 8- or 16-bit integer. Combined with a 16-bit weight, this yields roughly 24 bits of effective precision, analogous to the PXR24 format from high-dynamic-range imaging.

For optimizer states, they divide tensors into groups of 32 elements, normalize each group by its maximum absolute value, and apply a nonlinear function before linear quantization to 8-bit integers. Momentum uses a softsign-like function, φ(x) = 2x/(1+|x|), which pulls extreme values toward the center. Variance applies a square root, φ(x) = √x, motivated by Adam's accumulation of squared gradients and the resulting heavy-tailed distribution. Both are exactly invertible and require no hyperparameters. A separate FP16 scale per group adds 2/32 = 1/16 bytes per parameter of overhead. The whole step is fused into a single Triton kernel.

Evaluation uses SGD, AdamW, and Lion, with identical hyperparameters between reference optimizers and FlashOptim counterparts. Reference optimizers are re-implemented with similar fused Triton kernels. Reference configurations use FP32 for weights, gradients, momentum, and variance, with BF16 activations; FlashOptim uses BF16 weights, INT8 correction, BF16 gradients, INT8 momentum, UINT8 variance, and BF16 activations. Tasks include ResNet-50 on ImageNet-1K, GPT-2 (124M) pretraining on a 10B-token subset of FineWeb using the nanoGPT recipe, and supervised finetuning of Llama-3.1-8B on OpenMathInstruct-2 evaluated on GSM8k. Training used PyTorch 2.8 and CUDA 12.8 on NVIDIA H100 GPUs, with mean and standard deviation over 3 random seeds.

Why This Matters

Impact on research. FlashOptim reduces per-parameter memory by over 50% without changing optimizer semantics, meaning it can be applied as a drop-in replacement rather than requiring changes to training dynamics, as parameter-efficient methods do. The paper also shows that a simple companding function is enough to reach 8-bit optimizer states that were previously attributed to elaborate software complexity, and it suggests that designing custom companding functions is a fruitful direction for future work. The authors state that lowering memory requirements enables practitioners and researchers with limited hardware to train larger models than previously feasible.

Real-world applications:

• Finetuning large language models on limited numbers of GPUs, where optimizer state and master weights dominate memory.
• Image classification training pipelines using SGD or AdamW where activation memory is not the bottleneck.
• Distributed training with FSDP, where the reduced per-rank footprint composes with sharding for multiplicative savings.
• Checkpoint storage and transfer, where the smaller representation cuts checkpoint size by more than half (84 GiB to 35 GiB for a 7B model).

Industry relevance. The optimizations compose with existing memory-reduction infrastructure such as ZeRO, FSDP, CPU offloading, activation checkpointing, and freezing parameters. The memory savings are reported as effectively free in terms of throughput, and the implementation is released as an open-source PyTorch library, making adoption in existing training stacks straightforward.

Future Directions

• Custom companding functions. The authors' ablations across different tensor types suggest that designing custom companding functions is a fruitful direction for future research.
• Robustness across tasks and architectures. The authors note that some tasks and architectures may be more sensitive to quantization than those benchmarked, and that no quantization approach can guarantee quality preservation in all cases.
• Combination with other memory techniques. Since FlashOptim preserves optimizer semantics, it could be combined with communication compression techniques, parameter-efficient methods, and sharding or offloading approaches; the paper states it could be used alongside gradient communication compression.
• Precision limits. The authors state that while 24-bit master weights suffice in their experiments, not even 32 bits are guaranteed to suffice in all cases, because if the ratio of weight update to weight magnitude falls below the floating-point spacing, the update will be discarded.

Target Audience

This work is most useful to machine learning engineers and researchers who train or finetune large models under accelerator memory constraints, particularly those who want memory savings without altering optimizer semantics or adopting parameter-efficient alternatives. It is also relevant to systems researchers studying quantization, low-precision training numerics, and fused kernel implementations, and to practitioners deploying distributed training with FSDP or similar frameworks who want to reduce per-rank memory and checkpoint size. Readers without a background in floating-point representation and optimizer internals will need to consult the referenced prior work to follow the weight-splitting derivation.

Authors’ abstract

Standard mixed-precision training of neural networks requires many bytes of accelerator memory for each model parameter. These bytes reflect not just the parameter itself, but also its gradient and one or more optimizer state variables. With each of these values typically requiring 4 bytes, training even a 7 billion parameter model can be impractical for researchers with less than 100GB of accelerator memory. We introduce FlashOptim, a suite of optimizations that reduces per-parameter memory by over 50% while preserving model quality and API compatibility. Our approach introduces two key techniques. First, we improve master weight splitting by finding and exploiting a tight bound on its quantization error. Second, we design companding functions that greatly reduce the error in 8-bit optimizer state quantization. Together with 16-bit gradients, these techniques reduce AdamW memory from 16 bytes to 7 bytes per parameter, or 5 bytes with gradient release. They also cut model checkpoint sizes by more than half. Experiments with FlashOptim applied to SGD, AdamW, and Lion show no measurable quality degradation on any task from a collection of standard vision and language benchmarks, including Llama-3.1-8B finetuning.