Skip to content

Fix float16 zero-gradient NaNs in RMS transforms - #1785

Closed
vishesh9131 wants to merge 1 commit into
google-deepmind:mainfrom
vishesh9131:fix/float16-rms-stability
Closed

Fix float16 zero-gradient NaNs in RMS transforms#1785
vishesh9131 wants to merge 1 commit into
google-deepmind:mainfrom
vishesh9131:fix/float16-rms-stability

Conversation

@vishesh9131

@vishesh9131 vishesh9131 commented Sep 15, 2026

Copy link
Copy Markdown

Summary

  • compute RMS and centered-RMS denominators in at least float32
  • floor zero denominators and clamp centered variance against negative round-off
  • preserve the incoming gradient dtype after scaling
  • add JIT-compiled float16 regression coverage for both RMS transforms, epsilon placements, bias-correction modes, and scalar/ArrayLike epsilon values

Fixes #1784.

Tests

  • python -m optax._src.transform_test — 39 passed
  • python -m optax._src.alias_test — 592 passed, 62 skipped
  • pre-commit run --all-files inside the isolated test.sh environment — passed
  • bash test.sh proceeds through pre-commit, then stops at existing Pylint unexpected-keyword-arg errors for JAX out_sharding calls in optax/contrib/_sharding_test.py and optax/tree_utils/_tree_math.py; neither file is modified by this PR

@vishesh9131
vishesh9131 force-pushed the fix/float16-rms-stability branch from 668b4ad to 1f553aa Compare September 15, 2026 18:42
@rajasekharporeddy

rajasekharporeddy commented Oct 1, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the PR! LGTM.

@rajasekharporeddy
rajasekharporeddy self-requested a review October 1, 2026 16:28
@vishesh9131

vishesh9131 commented Oct 6, 2026

Copy link
Copy Markdown
Author

@rajasekharporeddy Thanks for taking a look! All checks are passing. Is there anything else you’d like me to address before this is ready for review and merge?

rajasekharporeddy
rajasekharporeddy reviewed Oct 7, 2026
View reviewed changes

@rajasekharporeddy rajasekharporeddy left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! @rdyro Please take a look at this PR.

@copybara-service copybara-service Bot mentioned this pull request Oct 22, 2026
Merged
@rajasekharporeddy

rajasekharporeddy commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator

Hi @vishesh9131, could you please rebase on to latest upstream main?

@vishesh9131
vishesh9131 force-pushed the fix/float16-rms-stability branch from 1f553aa to c8b76e0 Compare October 23, 2026 09:11
@vishesh9131

vishesh9131 commented Oct 23, 2026

Copy link
Copy Markdown
Author

done, you can merge it now

copybara-service Bot pushed a commit that referenced this pull request Oct 24, 2026
@rdyro
Raise a ValueError at trace time when eps underflows in the optimizer…
88af3ef 
… state dtype

Positive `eps` values smaller than the smallest normal number of the dtype they are added in (`jnp.finfo(dtype).tiny`) are flushed to zero. For example the default `eps=1e-8` is flushed to zero in float16 (smallest normal 2^-14 ~= 6.1e-5), which turns zero-gradient steps into `0 / 0 = NaN` (#1784, #1789).

Rather than silently changing the requested `eps`, this change:

- Adds `numerics.maybe_static_check_eps(eps, tree, name)`, which raises a `ValueError` when a concrete `eps` satisfies `0 < eps < jnp.finfo(dtype).tiny` for the dtype of `eps + leaf`. Traced `eps` (e.g. from `inject_hyperparams`) and `eps=0` are not checked.
- Calls it on `eps` (and `eps_root`) in every transform that adds `eps` to a per-parameter gradient-statistics accumulator:
  - `optax/_src/transform.py`: `scale_by_rss`, `scale_by_rms`, `scale_by_stddev`, `scale_by_adam`, `scale_by_amsgrad`, `scale_by_adamax`, `scale_by_adadelta`, `scale_by_adan`, `scale_by_belief`, `scale_by_yogi`, `scale_by_radam`, `scale_by_sm3`, `scale_by_novograd`. `optax.sm3` gains an `eps` argument so the check is actionable.
  - `optax/_src/factorized.py`: `scale_by_factored_rms` (`epsilon`, used by `adafactor`).
  - `optax/contrib`: `acprop`, `ademamix`, `simplified_ademamix`, `adopt`, `dadapt_adamw`, `madgrad`, `momo_adam`, `sophia`.
- Default `eps` values are unchanged. As a consequence, optimizers using these transforms with default `eps` now raise for float16 optimizer state; pass a larger `eps` (e.g. `eps=1e-4`) or keep the optimizer state in bfloat16/float32. float16 is removed from the dtype-stability test matrices and the sharding tests use bfloat16.

Thanks to @vishesh9131 (#1785) and @Anu27n (#1790) for reporting and the initial fixes.

PiperOrigin-RevId: 1002704127
copybara-service Bot pushed a commit that referenced this pull request Oct 24, 2026
@rdyro
Raise a ValueError at trace time when eps underflows in the optimizer…
336fa64 
… state dtype

Positive `eps` values smaller than the smallest normal number of the dtype they are added in (`jnp.finfo(dtype).tiny`) are flushed to zero. For example the default `eps=1e-8` is flushed to zero in float16 (smallest normal 2^-14 ~= 6.1e-5), which turns zero-gradient steps into `0 / 0 = NaN` (#1784, #1789).

Rather than silently changing the requested `eps`, this change:

- Adds `numerics.maybe_static_check_eps(eps, tree, name)`, which raises a `ValueError` when a concrete `eps` satisfies `0 < eps < jnp.finfo(dtype).tiny` for the dtype of `eps + leaf`. Traced `eps` (e.g. from `inject_hyperparams`) and `eps=0` are not checked.
- Calls it on `eps` (and `eps_root`) in every transform that adds `eps` to a per-parameter gradient-statistics accumulator:
  - `optax/_src/transform.py`: `scale_by_rss`, `scale_by_rms`, `scale_by_stddev`, `scale_by_adam`, `scale_by_amsgrad`, `scale_by_adamax`, `scale_by_adadelta`, `scale_by_adan`, `scale_by_belief`, `scale_by_yogi`, `scale_by_radam`, `scale_by_sm3`, `scale_by_novograd`. `optax.sm3` gains an `eps` argument so the check is actionable.
  - `optax/_src/factorized.py`: `scale_by_factored_rms` (`epsilon`, used by `adafactor`).
  - `optax/contrib`: `acprop`, `ademamix`, `simplified_ademamix`, `adopt`, `dadapt_adamw`, `madgrad`, `momo_adam`, `sophia`.
- Default `eps` values are unchanged. As a consequence, optimizers using these transforms with default `eps` now raise for float16 optimizer state; pass a larger `eps` (e.g. `eps=1e-4`) or keep the optimizer state in bfloat16/float32. float16 is removed from the dtype-stability test matrices and the sharding tests use bfloat16.

Thanks to @vishesh9131 (#1785) and @Anu27n (#1790) for reporting and the initial fixes.

PiperOrigin-RevId: 1002704127
copybara-service Bot pushed a commit that referenced this pull request Oct 24, 2026
@rdyro
Raise a ValueError at trace time when eps underflows in the optimizer…
9a1f892 
… state dtype

Positive `eps` values smaller than the smallest normal number of the dtype they are added in (`jnp.finfo(dtype).tiny`) are flushed to zero. For example the default `eps=1e-8` is flushed to zero in float16 (smallest normal 2^-14 ~= 6.1e-5), which turns zero-gradient steps into `0 / 0 = NaN` (#1784, #1789).

Rather than silently changing the requested `eps`, this change:

- Adds `numerics.maybe_static_check_eps(eps, tree, name)`, which raises a `ValueError` when a concrete `eps` satisfies `0 < eps < jnp.finfo(dtype).tiny` for the dtype of `eps + leaf`. Traced `eps` (e.g. from `inject_hyperparams`) and `eps=0` are not checked.
- Calls it on `eps` (and `eps_root`) in every transform that adds `eps` to a per-parameter gradient-statistics accumulator:
  - `optax/_src/transform.py`: `scale_by_rss`, `scale_by_rms`, `scale_by_stddev`, `scale_by_adam`, `scale_by_amsgrad`, `scale_by_adamax`, `scale_by_adadelta`, `scale_by_adan`, `scale_by_belief`, `scale_by_yogi`, `scale_by_radam`, `scale_by_sm3`, `scale_by_novograd`. `optax.sm3` gains an `eps` argument so the check is actionable.
  - `optax/_src/factorized.py`: `scale_by_factored_rms` (`epsilon`, used by `adafactor`).
  - `optax/contrib`: `acprop`, `ademamix`, `simplified_ademamix`, `adopt`, `dadapt_adamw`, `madgrad`, `momo_adam`, `sophia`.
- Default `eps` values are unchanged. As a consequence, optimizers using these transforms with default `eps` now raise for float16 optimizer state; pass a larger `eps` (e.g. `eps=1e-4`) or keep the optimizer state in bfloat16/float32. float16 is removed from the dtype-stability test matrices and the sharding tests use bfloat16.

Thanks to @vishesh9131 (#1785) and @Anu27n (#1790) for reporting and the initial fixes.

PiperOrigin-RevId: 1005981980
@rajasekharporeddy

rajasekharporeddy commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator

Thanks for reporting #1784, putting together this fix, and rebasing, @vishesh9131!

This float16 eps underflow issues has now been resolved repo-wide in #1835 (with credit to you and this PR in the commit message). Rather than silently promoting the RMS denominator to float32, Optax now checks eps at trace time via numerics.maybe_static_check_eps across all transforms and raises a ValueError when 0 < eps < jnp.finfo(dtype).tiny, prompting users to pass a larger eps (e.g., eps=1e-4) or use a higher-precision optimizer state dtype (bfloat16/float32).

Closing this PR (and #1784) as superseded by #1835. Thanks again for your contribution!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@rajasekharporeddy rajasekharporeddy rajasekharporeddy left review comments
@rdyro rdyro Awaiting requested review from rdyro

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

RMS transforms produce NaN on float16 zero-gradient steps

2 participants

@vishesh9131 @rajasekharporeddy