FlashOptim: Optimizers for Memory Efficient Training
| Authors | Jose Javier Gonzalez Ortiz et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.23349 |
| Download | |
| Categories | cs.LG, cs.AI |
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.
Engineering Breakdown
Plain English
FlashOptim is a suite of memory optimization techniques that reduces the per-parameter memory footprint during neural network training by over 50% while maintaining model quality and keeping the same API. Standard mixed-precision training stores four values per parameter (the parameter itself, its gradient, and optimizer state variables), each typically 4 bytes, making it prohibitively expensive to train large models like 7B parameter networks on GPUs with less than 100GB of memory. The paper introduces two key innovations: an improved master weight splitting technique that tightly bounds quantization error, and specialized companding functions (non-linear compression/expansion) that dramatically reduce errors when storing optimizer states in 8-bit precision. These techniques allow researchers with modest hardware constraints to train large models practically.
Core Technical Contribution
The core innovation lies in two complementary techniques that work together to reduce optimizer memory overhead. First, the authors provide a tight theoretical bound on the quantization error introduced by master weight splitting—a technique where you store weights in low precision but maintain a higher-precision copy for parameter updates. This bound allows them to quantize more aggressively while knowing exactly what quality loss to expect. Second, they introduce companding functions (similar to mu-law compression used in audio) that non-linearly compress optimizer state distributions (like momentum and variance in Adam) into 8-bit integers, exploiting the fact that these distributions are highly skewed and don't need uniform quantization. Together, these move beyond naive quantization approaches that treat all optimizer states equally, instead adapting compression to the actual statistics of what needs to be stored.
How It Works
The system operates on the standard training loop but intercepts memory allocation for optimizer state variables. During the forward pass, parameters are stored in low precision (likely 8-bit or 16-bit). For the backward pass, gradients are computed in standard precision but immediately compressed using the companding function—a learnable or pre-computed non-linear mapping that maps floating-point gradient statistics into 8-bit buckets more efficiently than linear quantization. The master weight copy (a higher-precision version of parameters for numerical stability) is maintained, but its quantization error is bounded by the theoretical analysis, allowing safe aggressive quantization. During parameter updates, the optimizer state (momentum buffers, variance estimates, etc.) is decompressed on-the-fly from 8-bit, the update is computed, and then immediately recompressed back to 8-bit. The entire process maintains numerical stability while reducing memory from roughly 4 bytes × 3 values per parameter to approximately 2 bytes × 3 values, a >50% reduction.
Production Impact
For teams training large language models or vision models on constrained hardware, this directly addresses a critical bottleneck: the memory wall that prevents training on consumer or mid-tier GPUs. A researcher with a single 40GB A100 GPU can now practically train models that previously required two or three GPUs, reducing infrastructure costs and iteration time. The approach maintains full API compatibility, meaning it can be integrated into existing PyTorch or JAX training loops with minimal code changes—likely a drop-in optimizer replacement. However, there are trade-offs: decompression/recompression of optimizer states adds modest compute overhead (roughly 5-10% slower training per step), and the companding functions require careful tuning per optimizer type (Adam, AdamW, SGD with momentum all have different state distributions). Integration requires instrumenting the backward pass and optimizer step, which may complicate debugging and performance profiling. For production serving, this has no impact since you only store final model weights.
Limitations and When Not to Use This
The paper assumes a fixed training setup with standard optimizers (likely Adam-family or SGD); it doesn't address other memory-hungry components like gradient checkpointing trade-offs or activation recomputation, which may be necessary for truly enormous models. The companding function design appears to require per-optimizer tuning and possibly per-model-architecture tuning (Transformer vs. CNN vs. RNN), which limits portability—a companding function optimal for BERT training may not be optimal for vision transformers. The paper doesn't clearly address numerical stability in edge cases or provide bounds on how error accumulates over thousands of training steps; there's a risk of subtle convergence issues that only appear after weeks of training. The approach is purely memory-centric and doesn't optimize for compute throughput or latency, so on hardware with abundant memory (like H100 clusters), it provides no benefit and may actually slow training. Finally, the abstract is incomplete, suggesting some key technical details on the companding function design are missing from what's provided.
Research Context
This work builds directly on prior research in quantization-aware training and mixed-precision methods, extending them to the optimizer state domain rather than just parameter or activation quantization. It relates to other memory-reduction techniques like LoRA and adapter methods, but takes a different angle—instead of reducing trainable parameters, it reduces memory per parameter through compression. The paper advances the line of work on training efficiency initiated by techniques like gradient accumulation and rematerialization (gradient checkpointing), offering an orthogonal dimension to explore. It's positioned for practical impact in the 7B-70B parameter model range where researchers frequently hit memory walls; this size range is particularly important as models move beyond research labs into smaller companies and academic institutions. The work likely opens follow-up directions in adaptive companding (per-layer or per-parameter-group compression), joint optimization with other memory techniques, and extension to distributed training scenarios where memory bottlenecks manifest differently.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::