Gate fp8 NaN gradient sanitization on quantization config

The NaN sanitization introduced for fp8 delayed-scaling FSDP ran on
every step regardless of quantization, adding a ~2-3% step-time
regression on non-fp8 workloads (per-float-grad jnp.nan_to_num
tree_map). The failure mode only occurs under fp8, so gate the block
on config.quantization in {"fp8", "fp8_full", "nanoo_fp8"}.

Non-fp8 workloads skip the tree_map entirely; fp8 behavior is
unchanged (verified: step 1 loss still finite under gpt3-52k + FSDP=8).

Author: ecnal-cienet

src/maxtext/trainers/pre_train/train.py

@@ -362,24 +362,26 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
   # values (skip the amax update for that step) instead of letting NaN flow through.
   # Also restore OWG values after apply_gradients to bypass optimizer corruption
   # (Adam should not update fp8 scale/amax_history).
-  fp8_stats = dict(grads).get(maxtext_utils.OVERWRITE_WITH_GRADIENT, None)
-  if fp8_stats is not None:
-    if maxtext_utils.OVERWRITE_WITH_GRADIENT in state.params:
-      current_fp8 = state.params[maxtext_utils.OVERWRITE_WITH_GRADIENT]
-      fp8_stats = jax.tree_util.tree_map(
-          lambda new, cur: jnp.where(jnp.isnan(new), cur, new),
-          fp8_stats,
-          current_fp8,
-      )
-    else:
-      fp8_stats = jax.tree_util.tree_map(lambda x: jnp.nan_to_num(x, nan=0.0), fp8_stats)
-    grads = dict(grads)
-    grads[maxtext_utils.OVERWRITE_WITH_GRADIENT] = fp8_stats
-  # Zero out any remaining NaN in float gradients to prevent param corruption
-  grads = jax.tree_util.tree_map(
-      lambda x: jnp.nan_to_num(x, nan=0.0) if jnp.issubdtype(x.dtype, jnp.floating) else x,
-      grads,
-  )
+  fp8_stats = None
+  if config.quantization in ("fp8", "fp8_full", "nanoo_fp8"):
+    fp8_stats = dict(grads).get(maxtext_utils.OVERWRITE_WITH_GRADIENT, None)
+    if fp8_stats is not None:
+      if maxtext_utils.OVERWRITE_WITH_GRADIENT in state.params:
+        current_fp8 = state.params[maxtext_utils.OVERWRITE_WITH_GRADIENT]
+        fp8_stats = jax.tree_util.tree_map(
+            lambda new, cur: jnp.where(jnp.isnan(new), cur, new),
+            fp8_stats,
+            current_fp8,
+        )
+      else:
+        fp8_stats = jax.tree_util.tree_map(lambda x: jnp.nan_to_num(x, nan=0.0), fp8_stats)
+      grads = dict(grads)
+      grads[maxtext_utils.OVERWRITE_WITH_GRADIENT] = fp8_stats
+    # Zero out any remaining NaN in float gradients to prevent param corruption
+    grads = jax.tree_util.tree_map(
+        lambda x: jnp.nan_to_num(x, nan=0.0) if jnp.issubdtype(x.dtype, jnp.floating) else x,
+        grads,
+    )
 
   if config.optimizer_memory_host_offload:
     state = state.replace(