Fix NaN in scale_by_adam on float16 zero-gradient steps

optax/_src/numerics.py

@@ -103,6 +103,48 @@ def safe_root_mean_squares(
   return jnp.where(rms <= min_rms, min_rms, jnp.sqrt(jnp.mean(abs_sq(x))))
 
 
+def add_eps_in_safe_dtype(
+    x: jax.typing.ArrayLike,
+    eps: jax.typing.ArrayLike,
+    min_dtype: jax.typing.DTypeLike = jnp.float32,
+) -> jax.Array:
+  """Returns ``x + eps``, computed in at least ``min_dtype`` precision.
+
+  ``eps`` is typically a small constant (e.g. the default ``1e-8`` used by
+  several optimizers) added to a value before taking its square root, to
+  keep a downstream division safe from a zero denominator. If ``x`` has a
+  low precision dtype (e.g. ``float16``, whose smallest representable
+  subnormal is ``~6e-8``), adding such an ``eps`` directly can silently
+  round to exactly ``x`` -- typically ``0.0`` when ``x`` is itself ``0``,
+  which defeats the purpose of the safety term and can turn it into a
+  ``0 / 0`` NaN. This function instead promotes ``x`` to at least
+  ``min_dtype`` before the addition, so ``eps`` is preserved.
+
+  This is a no-op change in behavior (bit-for-bit identical) when ``x`` is
+  already at least ``min_dtype`` precision (e.g. the default ``float32``
+  or ``float64``); only lower-precision dtypes (``float16``, ``bfloat16``)
+  are affected.
+
+  Callers that need the *result* of a computation built on top of this sum
+  back in ``x``'s original dtype (e.g. to keep an optimizer's returned
+  updates at the same precision as the incoming gradients) must cast back
+  explicitly -- this function only fixes the addition itself.
+
+  Args:
+    x: array the (possibly too-small-to-represent) ``eps`` is added to.
+    eps: term to add, typically a small non-negative constant.
+    min_dtype: the minimum floating point precision to compute the addition
+      in. Defaults to ``float32``.
+
+  Returns:
+    ``x.astype(safe_dtype) + eps``, where ``safe_dtype =
+    jnp.promote_types(x.dtype, min_dtype)``.
+  """
+  x = jnp.asarray(x)
+  safe_dtype = jnp.promote_types(x.dtype, min_dtype)
+  return x.astype(safe_dtype) + eps
+
+
 def safe_increment(count: jax.typing.ArrayLike) -> jax.Array:
   """Increments counter by one while avoiding overflow.
 

optax/_src/transform.py

@@ -136,11 +136,27 @@ def update_fn(updates, state, params=None):
     else:
       count_inc = jnp.asarray(0)
       nu_hat = nu
+    # `eps` is added in at least float32 precision (see
+    # `numerics.add_eps_in_safe_dtype`) so it does not silently underflow to
+    # zero under a low precision dtype like float16, which would otherwise
+    # turn `rsqrt(0)`/`1 / (0 + eps)` into `inf`, and then `inf * 0` (e.g. on
+    # a zero-gradient step) into a NaN update.
     if eps_in_sqrt:
-      scaling = jax.tree.map(lambda n: jax.lax.rsqrt(n + eps), nu_hat)
+      scaling = jax.tree.map(
+          lambda n: jax.lax.rsqrt(numerics.add_eps_in_safe_dtype(n, eps)),
+          nu_hat,
+      )
     else:
-      scaling = jax.tree.map(lambda n: 1 / (jnp.sqrt(n) + eps), nu_hat)
-    updates = jax.tree.map(lambda s, g: s * g, scaling, updates)
+      def _scaling(n):
+        denom = jnp.sqrt(numerics.add_eps_in_safe_dtype(n, 0.0)) + eps
+        return 1 / denom
+      scaling = jax.tree.map(_scaling, nu_hat)
+    # `scaling` is now at least float32 even if `updates`/`g` is float16;
+    # upcast `g` for the multiply, then cast the product back down to `g`'s
+    # original dtype so this transformation's output dtype is unchanged.
+    updates = jax.tree.map(
+        lambda s, g: (s * g.astype(s.dtype)).astype(g.dtype), scaling, updates
+    )
     if bias_correction:
       new_state = ScaleByRmsWithCountState(count=count_inc, nu=nu)
     else:
@@ -213,19 +229,28 @@ def update_fn(updates, state, params=None):
       mu_hat = mu
       nu_hat = nu
 
+    # See the matching comment in `scale_by_rms` -- `eps` is added in at
+    # least float32 precision so it does not silently underflow to zero
+    # under a low precision dtype like float16.
     if eps_in_sqrt:
       scaling = jax.tree.map(
-          lambda m, n: jax.lax.rsqrt(n - abs_sq(m) + eps),
+          lambda m, n: jax.lax.rsqrt(
+              numerics.add_eps_in_safe_dtype(n - abs_sq(m), eps)
+          ),
           mu_hat,
           nu_hat,
       )
     else:
-      scaling = jax.tree.map(
-          lambda m, n: 1 / (jnp.sqrt(n - abs_sq(m)) + eps),
-          mu_hat,
-          nu_hat,
-      )
-    updates = jax.tree.map(lambda s, g: s * g, scaling, updates)
+      def _scaling(m, n):
+        denom = jnp.sqrt(numerics.add_eps_in_safe_dtype(n - abs_sq(m), 0.0))
+        return 1 / (denom + eps)
+      scaling = jax.tree.map(_scaling, mu_hat, nu_hat)
+    # `scaling` is now at least float32 even if `updates`/`g` is float16;
+    # upcast `g` for the multiply, then cast the product back down to `g`'s
+    # original dtype so this transformation's output dtype is unchanged.
+    updates = jax.tree.map(
+        lambda s, g: (s * g.astype(s.dtype)).astype(g.dtype), scaling, updates
+    )
     if bias_correction:
       new_state = ScaleByRStdDevWithCountState(count=count_inc, mu=mu, nu=nu)
     else:
@@ -296,8 +321,21 @@ def update_fn(updates, state, params=None):
     # Algorithm 2 further multiplies Adam's standard nu_hat by b2. It is
     # unclear why. Other Nadam implementations also omit the extra b2 factor.
     nu_hat = optax.tree.bias_correction(nu, b2, count_inc)
+
+    def _update(m, v):
+      if m is None:
+        return None
+      # `eps_root`/`eps` are added in at least float32 precision (see
+      # `numerics.add_eps_in_safe_dtype`) so they do not silently underflow
+      # to zero when `v` has a low precision dtype (e.g. float16, whose
+      # smallest subnormal is ~6e-8). Otherwise the denominator can become
+      # exactly 0, turning this safety term into a 0/0 NaN whenever `m` and
+      # `v` are also 0 (e.g. on a zero-gradient step).
+      denom = jnp.sqrt(numerics.add_eps_in_safe_dtype(v, eps_root)) + eps
+      return (m.astype(denom.dtype) / denom).astype(m.dtype)
+
     updates = jax.tree.map(
-        lambda m, v: None if m is None else m / (jnp.sqrt(v + eps_root) + eps),
+        _update,
         mu_hat,
         nu_hat,
         is_leaf=lambda x: x is None,
@@ -376,8 +414,18 @@ def update_fn(updates, state, params=None):
       nu_eff = nu
 
     nu_max = jax.tree.map(jnp.maximum, state.nu_max, nu_eff)
+
+    def _update(m, v):
+      if m is None:
+        return None
+      # See the matching comment in `scale_by_adam` -- `eps_root`/`eps` are
+      # added in at least float32 precision so they do not silently
+      # underflow to zero under a low precision dtype like float16.
+      denom = jnp.sqrt(numerics.add_eps_in_safe_dtype(v, eps_root)) + eps
+      return (m.astype(denom.dtype) / denom).astype(m.dtype)
+
     updates = jax.tree.map(
-        lambda m, v: None if m is None else m / (jnp.sqrt(v + eps_root) + eps),
+        _update,
         mu_hat,
         nu_max,
         is_leaf=lambda x: x is None,
@@ -745,8 +793,20 @@ def update_fn(updates, state, params=None):
     else:
       mu_hat = optax.tree.bias_correction(mu, b1, count_inc)
     nu_hat = optax.tree.bias_correction(nu, b2, count_inc)
+
+    def _update(m, v):
+      if m is None:
+        return None
+      # See the matching comment in `scale_by_adam` -- `eps` is added in at
+      # least float32 precision so it does not silently underflow to zero
+      # under a low precision dtype like float16. `eps_root` was already
+      # folded into `v` above (before bias correction), so it needs no
+      # separate handling here.
+      denom = jnp.sqrt(numerics.add_eps_in_safe_dtype(v, 0.0)) + eps
+      return (m.astype(denom.dtype) / denom).astype(m.dtype)
+
     updates = jax.tree.map(
-        lambda m, v: None if m is None else m / (jnp.sqrt(v) + eps),
+        _update,
         mu_hat,
         nu_hat,
         is_leaf=lambda x: x is None,

optax/_src/transform_test.py

@@ -16,6 +16,8 @@
 
 """Tests of gradient transformations."""
 
+import functools
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -73,6 +75,50 @@ def test_scalers(self, scaler_constr):
     test_utils.assert_tree_all_finite((params, updates, state))
     test_utils.assert_trees_all_equal_shapes(params, updates)
 
+  @parameterized.named_parameters([
+      ('adam', transform.scale_by_adam),
+      ('amsgrad', transform.scale_by_amsgrad),
+      ('belief', transform.scale_by_belief),
+      (
+          'rms_eps_in_sqrt',
+          functools.partial(transform.scale_by_rms, eps_in_sqrt=True),
+      ),
+      (
+          'rms_eps_outside_sqrt',
+          functools.partial(transform.scale_by_rms, eps_in_sqrt=False),
+      ),
+      (
+          'stddev_eps_in_sqrt',
+          functools.partial(transform.scale_by_stddev, eps_in_sqrt=True),
+      ),
+      (
+          'stddev_eps_outside_sqrt',
+          functools.partial(transform.scale_by_stddev, eps_in_sqrt=False),
+      ),
+  ])
+  def test_no_nan_on_float16_zero_grad(self, scaler_constr):
+    """A zero-gradient float16 step must not produce NaN updates.
+
+    `eps` (and, where applicable, `eps_root`) default to values like `1e-8`,
+    which round to exactly 0.0 in float16 (whose smallest subnormal is
+    ~6e-8). That used to turn each of these transformations' `sqrt(v) + eps`
+    -style safety denominator into a hard 0, so `m / 0` (or `1 / 0`, then
+    `0 * grad`) produced NaN whenever the moment estimates were also 0, as
+    on this all-zero-gradient step. Every scaler here shares the same
+    underlying `... + eps` pattern and used to fail the same way.
+    """
+    params = jnp.array([1.0, -2.0, 0.5], dtype=jnp.float16)
+    zero_grads = jnp.zeros_like(params)
+    scaler = scaler_constr()
+    state = scaler.init(params)
+    for _ in range(3):
+      updates, state = scaler.update(zero_grads, state, params)
+      test_utils.assert_tree_all_finite(updates)
+      # The fix must not silently change this transformation's output dtype
+      # -- only the (still float16) NaN it used to produce.
+      test_utils.assert_trees_all_equal_dtypes(updates, zero_grads)
+      params = update.apply_updates(params, updates)
+
   def test_apply_every(self):
     # The frequency of the application of sgd
     k = 4