FlashOptim: Breaking the Memory Wall with 24-bit Weights and Companded States

TL;DR

Training a 7B parameter model shouldn't require 100GB+ of VRAM just for the optimizer. FlashOptim is a new suite of optimizers that slashes parameter-associated memory by over 50% (from 16 bytes/param to 7 bytes/param for AdamW) without losing a single point of accuracy. It achieves this by rethinking how we store "master weights" and by using a clever mathematical trick called "companding" to squeeze optimizer states into 8 bits.

The Problem: The Hidden Cost of "Master Weights"

In modern mixed-precision training, we perform forward and backward passes in 16-bit (BF16/FP16) for speed, but we keep a "Master" copy of the weights in FP32 to ensure small gradient updates aren't lost to rounding. When you add up the Master Weights (4 bytes), Gradients (4 bytes), and Adam states (8 bytes), you end up with 16 bytes per parameter.

For a Llama-3-8B model, that’s 128GB of VRAM before you even count the activations! FlashOptim identifies that most of these bits are redundant or poorly utilized.

Methodology: The Secret Sauce

FlashOptim introduces two surgical strikes against memory waste:

1. ULP-based Weight Splitting

Instead of a full 32-bit float, FlashOptim uses a 16-bit base weight ($het{a}^{\prime }$) and an 8-bit "correction" term ($\rho$).

The intuition? The error between a 32-bit float and its 16-bit version is bounded by the Unit in Last Place (ULP). By normalizing the error to this tiny interval before quantizing, FlashOptim achieves 24-bit effective precision. This is virtually indistinguishable from FP32 for training dynamics but saves 1 byte per parameter.

2. Companded Quantization

Standard 8-bit quantization assumes values are spread evenly. However, optimizer moments (momentum and variance) are often "heavy-tailed"—most values are tiny, but a few are huge.

FlashOptim applies a companding function (Compression + Expanding) before quantization:

• Momentum: Uses a softsign-like function to spread values more evenly.
• Variance: Uses a square root to squash the massive dynamic range of squared gradients.

Visual Evidence: Companding (orange) reshapes the distribution of optimizer states to better utilize the 256 available bins in INT8/UINT8.

Performance: SOTA Results, Half the Memory

The authors tested FlashOptim across ImageNet (ResNet-50) and LLM Pretraining/Finetuning (GPT-2, Llama-3.1-8B).

• Convergence: FlashOptim curves are pixel-perfect matches for the baseline.
• Memory: In Llama-3.1-8B finetuning, peak memory dropped from 175 GiB to 113 GiB.
• Speed: Since the logic is fused into high-performance Triton kernels, the overhead is non-existent.

Figure: Memory saving breakdown for Llama-3.1-8B.

Critical Insight: Why Companding Matters

One of the most striking findings is that linear quantization destroys training. Without the non-linear companding step, GPT-2 training diverges almost instantly when using 8-bit variance. By reshaping the data distribution, FlashOptim makes low-precision training numerically stable.

Figure: Without companding (red), training diverges. With companding (blue), it matches FP32.

Conclusion & Perspective

FlashOptim is a "free lunch" for AI researchers. It provides drop-in replacements for torch.optim that significantly lower the hardware barrier. While it won't help much for small models with huge activation maps (like high-res CV), for the current era of Transformer-heavy LLMs, it is an essential tool for maximizing GPU efficiency.

Future Outlook: Could we go even lower? The success of 8-bit companded states suggests that 4-bit optimizer states might be feasible if we find even more aggressive non-linear transformations.