diff --git a/optax/_src/numerics.py b/optax/_src/numerics.py index fca01df12..4327a9220 100644 --- a/optax/_src/numerics.py +++ b/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. diff --git a/optax/_src/transform.py b/optax/_src/transform.py index ec2e8a2b5..9499275fb 100644 --- a/optax/_src/transform.py +++ b/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, diff --git a/optax/_src/transform_test.py b/optax/_src/transform_test.py index f521045ff..f30d2c433 100644 --- a/optax/_src/transform_test.py +++ b/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