fix: update

Author: mesakhcienet

src/maxtext/layers/linears.py

@@ -94,7 +94,10 @@ def _compute_dot_general_nnx(
   if quant_dot_general is not None:
     if initializing:
       quant_dot_general.lazy_init(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None)
-    return quant_dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None, mutable=["aqt"])
+    return quant_dot_general(
+        inputs, kernel, ((axis, contract_ind), ((), ())),
+        precision=None, mutable=["aqt", "_overwrite_with_gradient"],
+    )
 
   return dot_general(
       inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision, out_sharding=out_sharding

src/maxtext/layers/nnx_decoders.py

@@ -530,9 +530,8 @@ def pure_layer_fn(state_in, y_in):
       out = merged_layer(y_in, **kwargs)
       return out, nnx.state(merged_layer)
 
-    if not self._has_linen_fp8_side_effects():
-      pure_layer_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
-    out, new_state = pure_layer_fn(state, y)
+    checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
+    out, new_state = checkpointed_fn(state, y)
     nnx.update(layer, new_state)
 
     return out
@@ -574,8 +573,7 @@ def layer_fn(carry, scanned_vars):
       # ONLY return non-param state to prevent memory duplication of weights
       return new_carry, new_current_state
 
-    if not self._has_linen_fp8_side_effects():
-      layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+    layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
 
     final_carry, scanned_other = jax.lax.scan(layer_fn, x_in, (params, state))
 
@@ -647,19 +645,6 @@ def minimal_policy(self, with_context=False, with_quantization=False):
       names.append("quantization")
     return jax.checkpoint_policies.save_only_these_names(*names)
 
-  def _has_linen_fp8_side_effects(self):
-    """Check if the current quantization uses Linen FP8 modules that create mutable state.
-
-    FP8 GPU/NANOO quantization with QWIX creates Linen FP8 modules (e.g.,
-    nn.Fp8DirectDotGeneralOp, nn.NANOOFp8DotGeneralOp) during the forward pass.
-    These modules use self.variable() to create mutable state (amax histories,
-    scales) as side effects. When called inside jax.checkpoint, these side effects
-    cause UnexpectedTracerError because the traced values escape the checkpoint scope
-    through the Linen variable scope.
-    """
-    cfg = self.config
-    return cfg.use_qwix_quantization and cfg.quantization in ("fp8_gpu", "fp8_nanoo")
-
   def get_remat_policy(self):
     """Get remat policy for jax.checkpoint."""
     policy = None
@@ -1154,10 +1139,7 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in):
         out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs)
         return out_y, out_kv, nnx.state(merged_layer)
 
-      if not self._has_linen_fp8_side_effects():
-        checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
-      else:
-        checkpointed_fn = pure_layer_fn
+      checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
 
       for lyr, layer in enumerate(self.layers):
         graphdef, state = nnx.split(layer)
@@ -1261,10 +1243,7 @@ def pure_gemma_fn(graphdef, state_in, y_in):
         )
         return out_y, nnx.state(merged_layer)
 
-      if not self._has_linen_fp8_side_effects():
-        checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse)
-      else:
-        checkpointed_gemma_fn = pure_gemma_fn
+      checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse)
 
       graphdef, state = nnx.split(self.layers_remainder)
       y, new_state = checkpointed_gemma_fn(graphdef, state, y)

src/maxtext/layers/quantizations.py

@@ -638,6 +638,10 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"):
     )
 
   if config.use_qwix_quantization:
+    if config.quantization == "fp8_gpu":
+      return Fp8Quantization()
+    if config.quantization == "fp8_nanoo":
+      return NANOOFp8Quantization()
     return None
   quant_cfg = _get_quant_config(config)
   if quant_cfg:
@@ -819,6 +823,10 @@ def maybe_quantize_model(model, config):
   """Quantize the model if quantization is enabled."""
   # Batch split is not using Qwix's interception feature but manual plumbing
   if config.use_qwix_quantization and not config.use_batch_split_schedule:
+    # fp8_gpu/fp8_nanoo dot_general is handled by DenseGeneral's ToNNX wrapper,
+    # bypassing QWIX interception to avoid tracer leaks inside jax.checkpoint.
+    if config.quantization in {"fp8_gpu", "fp8_nanoo"}:
+      return model
     quantization_provider = get_qt_provider(config)
     if quantization_provider:
       model = qwix.quantize_model(model, quantization_provider)