FlashOptim: drop-in memory-efficient optimizer replacements (Adam/SGD/Lion) via weight-splitting + ECC + companded 8-bit state

TL;DR

FlashOptim is a Databricks/Mosaic library of drop-in PyTorch optimizers (FlashSGD, FlashAdamW, FlashLion) that compute the same updates as their fp32 counterparts but cut peak training memory ~35% and checkpoint size ~57% on 8B-scale finetuning. Three ingredients: (1) 24-bit master weights stored as a low-precision view + an N-bit error-correction code that exactly recovers the fp32 value via interval-based integer quantization, (2) 8-bit + 16-bit-scale-per-32-elements quantized optimizer states with companding (invertible nonlinear preprocessing borrowed from signal processing) to make momentum/variance distributions fit the quantization grid without divergence, and (3) gradients kept at the backward-pass dtype with optional gradient release. Fused Triton kernels mean no practical step-time overhead. Compatible with DDP and FSDP2 (not FSDP1).

Key claims

The optimizer-state bytes are only one of three sources of training memory; ignoring master-weight duplication and gradient upcasting leaves most of the savings on the table — so any “memory-efficient optimizer” worth the name has to address all three [Thread 2/n, 14/16].
The downcast-redundancy fix is weight splitting: instead of carrying both an fp32 master and a bf16 view, store one B-bit downcast value + a (32-B)-bit ECC that exactly reconstructs the fp32 value, yielding 24-bit master weights from a 16-bit downcast + 8-bit ECC [Thread 4/n–5/n; README].
The novel ECC algorithm exploits that any downcast value x’ has a tight, computable interval [val_min, val_max] of fp32 values that map to it; the ECC is then just an N-bit integer-quantized position within that interval — yielding exact reconstruction for most bitstrings with N+B=32 and graceful degradation otherwise [Thread 6/n–8/n].
Naive 8-bit linear quantization (+ 16-bit scale per 32 elements) of Adam’s momentum and variance diverges out of the box; the fix is companding — different invertible preprocessing functions for momentum-like vs. variance-like state — which can flip B-bit quantization “from broken to great with one line of preprocessing” [Thread 9/n–12/n].
Companding is described as “a decades-old idea ubiquitous in compression and signal processing, but disturbingly absent from the deep learning literature”; the author argues tuned companding should be standard in DL model-compression papers rather than chasing new datatypes [Thread 10/n, 12/n].
Reported headline: finetuning an 8B model with FlashOptim uses 35% less peak memory and writes checkpoints 57% smaller than the fp32 baseline, without affecting model convergence [README].
All compression operations are fused into a single update Triton kernel, so the only step-time overhead is one-time JIT compilation on the first step [README].
Checkpoint serialization has a hidden gotcha: PyTorch’s Optimizer.load_state_dict() casts fp state to the parameter dtype (e.g. fp32→bf16), which is lossy; FlashOptim works around this by either serializing int8+fp16 scales directly (compress_state_dict=True, the default) or by pre-quantizing fp32 state to int8+scales before PyTorch’s cast runs (compress_state_dict=False) [README].
Compatibility: works with DDP and FSDP2 (fully_shard); explicitly does not support FSDP1 due to how FSDP1 manages parameter/optimizer-state sharding [README].
Supports gradient release — params are updated during the backward pass as soon as their gradient is computed, freeing gradient memory before the full backward completes [Thread 13/16; README].

Method

The training step is restructured around three quantized memory primitives. Master weights become (B bit downcast view) + (32−B bit ECC code), with default (B, 32−B) = (16, 8) → 24-bit semantics. The ECC is not the low bits of the fp32 value; it’s an integer-quantized location within the fp32 interval [val_min, val_max] that maps to the downcast view. This makes the splitting work across arbitrary downcast dtypes (not just bf16/fp32) and degrades smoothly when N < 24. Optimizer state (Adam’s m, v) is held at 8 bits with one fp16 scale per 32 elements — but only after passing through a per-state-type companding function, an invertible nonlinear preprocessor (analogous to µ-law companding in audio) that reshapes the value distribution to better match the uniform quantizer’s dynamic range. Without companding, plain 8-bit linear quantization fails to train; with it, the quantization error is small enough that updates are numerically equivalent to fp32 in convergence. Gradients stay in the backward-pass dtype (e.g. bf16) instead of being upcast to fp32, and optionally gradient release fuses the parameter step into the backward pass to reclaim gradient buffers eagerly. Everything is implemented as fused Triton kernels — compression, decompression, ECC reconstruction, companding, and the Adam/SGD/Lion update all live in one kernel per optimizer, with no per-step Python overhead.

Results

8B finetuning: 35% less peak training memory, 57% smaller checkpoints; convergence matches the fp32 baseline [README §Memory Savings].
Per-parameter memory: ~50% reduction for typical configs (cast model to bf16 + master_weight_bits=24 + 8-bit optimizer state) [README §Quick Start].
Step-time overhead: zero in practice after first-step Triton JIT compilation [README §Fused Triton Kernels].
Datatype coverage: FlashSGD, FlashSGDW, FlashAdam, FlashAdamW, FlashLion, all behind the standard PyTorch optimizer API.
Parallelism support: DDP and FSDP2 (fully_shard); FSDP1 unsupported by design.
Code is a single Python file plus the Triton kernels — the project deliberately optimizes for “you can pip install this and drop it into an existing training loop.”

Why it’s interesting

FlashOptim is the systems-side counterpart to two things the wiki has been tracking. (1) veScale-FSDP: Flexible and High-Performance FSDP at Scale argued the missing FSDP primitive is block-wise sharding granularity that natively supports 8-bit Adam and Muon without intrusive model code; FlashOptim is precisely the kind of 8-bit-Adam implementation that consumes that primitive, and it explicitly targets FSDP2 (the new sharding format) rather than FSDP1. The two stories compose: veScale-FSDP fixes how the shards are laid out, FlashOptim fixes what’s in each shard. (2) The “remove a hand-tuned axis” pattern from Training stability at scale — MuonClip, Quantile Balancing, SSO’s spectral retraction — extends here in spirit: companding eliminates the “tune your custom datatype” axis that prior 8-bit optimizers (bitsandbytes, AnyPrecision) leaned on, replacing it with a one-line preprocessing function that makes vanilla 8-bit quantization work. Worth checking against Preconditioned Inexact Stochastic ADMM for Deep Models (PISA / SISA / NSISA)‘s PISA claim that ADMM-based optimizers beat AdamW/Muon/Shampoo — FlashOptim does not change the update math, only its memory footprint, so the two papers occupy orthogonal axes and could in principle be stacked.