[Feat]: Stochastic rounding of AdamW momentums

[kohya-ss]

Describe your use-case.

With Adam/AdamW, gradients and optimizer state appear to stay in bf16, which can lead to a loss of precision during the accumulation of gradients over many steps, potentially impacting training stability and final model quality. Adafactor already does the math/state in fp32 and only stochastic rounding on copy-back.

This proposal aims to align the Adam optimizer with this best practice, which should improve its performance in bfloat16 training scenarios.

What would you like to see as a solution?

I suggest modifying step_adam_parameter in modules/util/optimizer/adam_extensions.py (and adamw_extensions.py) to ensure internal calculations are performed in float32.

A simple way to achieve this would be to cast the gradient to float32 before it is used for state updates, similar to how it is handled in the Adafactor implementation.

Have you considered alternatives? List them here.

• Keep current bf16-state behavior to minimize VRAM (at some accuracy cost).
• Expose a flag, e.g. state_dtype = "fp32" | "param", defaulting to "fp32", so users can choose.
• Related: [Bug]: No weight decay with Adam, bf16 and stochastic rounding #996 (decay effectiveness with bf16 + stochastic rounding); see also [Bug]: Stochastic rounding in Adafactor is overridden #994 for stochastic rouning behavior in Adafactor.

[Koratahiu]

I don't think this is true; storing the states in bf16 shouldn't harm performance. I think Chroma is already trained with full bf16 (optimizer states, weights, etc.).

Also, this paper about bf16 training:
https://arxiv.org/abs/2502.20566
States that storing states in bf16 with SR should be fine and can do the same quality as full fp32 training.
The only thing to consider is that the training with full bf16 with SR requires x2-x5 higher LR to convergence at the same speed as fp32.

Adafactor stores its states in fp32 because its states are already factored and low-rank, so reducing its precision to bf16 isn't a good idea.

edit1:
Another source is this paper (the paper that shed light on SR for BF16):
https://arxiv.org/abs/2010.06192
Which used BF16 on all components (activations, gradients, model weights, and optimizer states) and achieved FP32 training quality by employing stochastic rounding (SR)..

[dxqb]

Thanks for looking at those details!

I suggest modifying step_adam_parameter in modules/util/optimizer/adam_extensions.py (and adamw_extensions.py) to ensure internal calculations are performed in float32.

I don't think it is possible to calculate in fp32 without storing momentums in fp32. This is easiest to see here in Adam:

OneTrainer/modules/util/optimizer/adamw_extensions.py

Line 144 in ea3bc37

addcdiv_stochastic_(p, exp_avg, denom, value=-step_size)

exp_avg is the bf16 momentum

Calculating and storing in fp32 might be beneficial, but who could use this functionality?
For full finetuning, nobody has the vram to store two momentums in fp32. It defeats the purpose of using bf16.
For LoRA, keeping parameters in fp32 is highly recommended - so we don't have any bf16 here.

#994 is fixed. #996 is an issue and should be fixed.

[dxqb]

Adding one thing:
Those calculations, that can be done in fp32 without storing in fp32, are already done in fp32. Those are hidden behind addcdiv_stochastic_. The addition and the divison are done in fp32 by this function.

[kohya-ss]

Thank you for the detailed explanation. I completely understand the constraints regarding VRAM, and I agree that memory efficiency is crucial for full finetuning.

Regarding LoRA, I also agree with your point that keeping parameters in float32 is recommended. I brought this up because I noticed that several LoRA presets specify "train_dtype": "BFLOAT_16", which led me to believe that stochastic rounding might be utilized in those cases as well.

My primary concern lies in the calculations that occur before the addcdiv_stochastic_ call, specifically with exp_avg.lerp_ and exp_avg_sq.addcmul_. Since exp_avg and grad are both in bfloat16, it appears these operations are performed in bfloat16 precision. I am concerned that this could lead to a loss of precision during the momentum accumulation stage.

This could be an intentional trade-off. But if it was unintended, there might be room to revisit the numerical precision in that part of the update.

Please let me know if I have misunderstood anything. Thank you again for your time and consideration.

[dxqb]

Regarding LoRA, I also agree with your point that keeping parameters in float32 is recommended. I brought this up because I noticed that several LoRA presets specify "train_dtype": "BFLOAT_16", which led me to believe that stochastic rounding might be utilized in those cases as well.

that's only used for activations and calculations. the parameter and gradient dtype is lora_weight_dtype

My primary concern lies in the calculations that occur before the addcdiv_stochastic_ call, specifically with exp_avg.lerp_ and exp_avg_sq.addcmul_. Since exp_avg and grad are both in bfloat16, it appears these operations are performed in bfloat16 precision. I am concerned that this could lead to a loss of precision during the momentum accumulation stage.

you could probably do some calculations there in fp32, but always cast back to bf16 each step. and even apply stochastic rounding to the momentum updates?

How much that'll help in practice I don't know.

stochastic rounding is applied quite incompletely in general: only for parameter updates. it could be used at various other points. we've recently discussed applying it to EMA because there are likely precision issues.

Please let me know if I have misunderstood anything. Thank you again for your time and consideration.

very welcome

[kohya-ss]

Thank you for clarifying the LoRA settings. I appreciate you taking the time to answer my question.

you could probably do some calculations there in fp32, but always cast back to bf16 each step. and even apply stochastic rounding to the momentum updates?

Yes, that's exactly the idea I had in mind—applying stochastic rounding to the optimizer states as well. I thought it could be a potential alternative to improve precision with minimal overhead, especially if holding the states in float32 is not feasible.

I'm glad if this discussion has been helpful in any way.

[dxqb]

Thank you for clarifying the LoRA settings. I appreciate you taking the time to answer my question.

you could probably do some calculations there in fp32, but always cast back to bf16 each step. and even apply stochastic rounding to the momentum updates?

Yes, that's exactly the idea I had in mind—applying stochastic rounding to the optimizer states as well. I thought it could be a potential alternative to improve precision with minimal overhead, especially if holding the states in float32 is not feasible.

in case you are thinking about implementing that, consider doing it as a PR to @Koratahiu 's optimizer. He has been implementing optimizers including AdamW with advanced functions such as stochastic rounding (and more)

ideally, we can stop the monkey-patching of upstream optimizers at some point (by having a good upstream that has everything we need).

repo is here: https://github.com/Koratahiu/Advanced_Optimizers
#992

documentation focuses on 1-bit because that's the main innovation, but it can be used non-factored.

[dxqb]

Koratahiu/Advanced_Optimizers#20