Repository navigation
Stochastic rounding support for 8-bit optimizers #1165
Description
Feature request
It would be extremely helpful to have stochastic rounding support, as described in Revisiting BFloat16 Training, for 8-bit optimizers. This would require working with a full-precision set of model updates and a half-precision set of model parameters, and to randomly round up or down on the parameter update with a chance determined by the quantization error. If implemented correctly, this should have minimal additional overhead and should perform similarly to full-precision training.
Motivation
This would allow even more memory savings, where one could fit the entire set of model parameters and an 8bit optimizer in the same space as one full-precision copy of model parameters, or less than that if using 8bit Lion.
Your contribution
Unfortunately, I don't know CUDA, and my understanding is that that would be required to implement this, since optimizer gradient transforms and parameter updates are all handled outside of Python in one inseparable process.
Activity
SirTrippsalot commentedon Apr 29, 2024
I'd love to see this added. Greatly improves BF16 training results on other optimizers.
wassname commentedon Jan 27, 2025
example implementation here https://github.com/warner-benjamin/optimi
TimDettmers commentedon Mar 7, 2025
I actually tried this in experiments and it did not help. We are happy to bring this back, the implementation is not that complicated.
Right now we are a bit short on people hours for features like this and outside contributions would be very welcome for this.
is-Xiaoen commentedon Oct 4, 2026
Would a bounded, experiment-first contribution still be useful here? Given the earlier finding that stochastic rounding did not help, the proposed first scope is opt-in stochastic rounding for BF16 parameter updates in non-paged CUDA AdamW8bit, keeping the existing 8-bit state quantization and avoiding persistent FP32 master weights.
The evaluation would compare nearest rounding and stochastic rounding against a higher-precision reference, with accumulated-update tests, several seeds, memory use, and optimizer-step overhead. Negative results would be reported before proposing a feature PR. Two scope details need agreement: should the initial implementation cover the FP32-state fallback for small tensors, and which weight-decay baseline should it use given #2040? It would also need to stay compatible with the stream work in #2050.
Is that a useful starting scope, and are there particular workloads or acceptance criteria you would want tested?