update

Author: mesakhcienet

src/maxtext/layers/nnx_decoders.py

@@ -527,14 +527,14 @@ def layer_fn(carry, scanned_vars):
     # with jax.checkpoint's purity requirement, causing UnexpectedTracerError.
     # Linen's nn.scan avoids this via variable_axes={"aqt": 0,
     # "_overwrite_with_gradient": 0}, but jax.lax.scan has no collection awareness.
-    # Only skip checkpoint for FP8-based quant types that have this tracer leak.
-    # Other quant types (AQT int8, Qwix) don't use mutable Linen state and are
-    # compatible with jax.checkpoint.
-    _fp8_quant_types = (quantizations.Fp8Quantization, quantizations.NANOOFp8Quantization)
-    _has_fp8_tracer_issue = isinstance(self.quant, _fp8_quant_types) or (
-        hasattr(quantizations, "TransformerEngineQuantization")
-        and isinstance(self.quant, quantizations.TransformerEngineQuantization)
-    )
+    # Skip checkpoint for FP8-based quantization configs. We check the config
+    # string rather than self.quant because qwix quantization (use_qwix_quantization=True)
+    # sets self.quant=None but still intercepts lax.dot_general at runtime with
+    # Linen FP8 modules (NvidaFp8Provider → nn.Fp8DirectDotGeneralOp).
+    _quant_str = getattr(self.config, "quantization", "") or ""
+    _has_fp8_tracer_issue = any(
+        fp8_id in _quant_str for fp8_id in ("fp8", "nanoo_fp8")
+    ) or _quant_str.startswith("te_")
     if not _has_fp8_tracer_issue:
       layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)