fix: update

Author: mesakhcienet

src/maxtext/layers/nnx_decoders.py

@@ -42,7 +42,6 @@
 from maxtext.layers import initializers, linears, mhc, normalizations, quantizations
 from maxtext.layers.attentions import Attention
 from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding
-from maxtext.layers.engram import Engram, NgramHashMapping
 from maxtext.layers.normalizations import RMSNorm
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.models import (
@@ -333,7 +332,7 @@ def __init__(
                   dense_cls, length=(next_boundary - current_idx), metadata_axis_name=chunk_name, rngs=rngs
               ))
               current_idx = next_boundary
-              
+
           # 2. Create MoE Chunks (Direct setattr, NO nnx.Dict)
           current_idx = config.first_num_dense_layers
           while current_idx < config.num_decoder_layers:
@@ -531,8 +530,9 @@ def pure_layer_fn(state_in, y_in):
       out = merged_layer(y_in, **kwargs)
       return out, nnx.state(merged_layer)
 
-    checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
-    out, new_state = checkpointed_fn(state, y)
+    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)
     nnx.update(layer, new_state)
 
     return out
@@ -574,7 +574,8 @@ def layer_fn(carry, scanned_vars):
       # ONLY return non-param state to prevent memory duplication of weights
       return new_carry, new_current_state
 
-    layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+    if not self._has_linen_fp8_side_effects():
+      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))
 
@@ -646,6 +647,19 @@ 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
@@ -935,7 +949,7 @@ def _find_next_boundary(self, current_idx, end_idx, engram_indices):
   def _apply_single_engram_layer(self, y, layer_name, *args, **kwargs):
     """Applies a single, unscanned Engram layer."""
     layer = getattr(self, layer_name)
-    
+
     decoder_input_tokens = kwargs.get("decoder_input_tokens")
     layer_kwargs = kwargs.get("layer_kwargs", {})
 
@@ -1000,7 +1014,7 @@ def _apply_interleaved_scanned_layers(self, y, layer_prefix, start_idx, end_idx,
         chunk_name = f"{layer_prefix}_{current_idx}_{next_boundary - 1}"
         chunk_stack = getattr(self, chunk_name)
         scan_length = next_boundary - current_idx
-        
+
         y, chunk_stack = self._apply_layers_sequentially(
             chunk_stack, y, *args, length=scan_length, **kwargs.get("layer_kwargs", {})
         )
@@ -1079,7 +1093,8 @@ def __call__(
           )
 
           y = self._apply_interleaved_scanned_layers(
-              y, "moe_layers", cfg.first_num_dense_layers, cfg.num_decoder_layers, cfg.engram_layers, *layer_args, **common_kwargs
+              y, "moe_layers", cfg.first_num_dense_layers, cfg.num_decoder_layers,
+              cfg.engram_layers, *layer_args, **common_kwargs
           )
         else:
           y, self.dense_layers = self._apply_layers_sequentially(
@@ -1139,7 +1154,10 @@ 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)
 
-      checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
+      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
 
       for lyr, layer in enumerate(self.layers):
         graphdef, state = nnx.split(layer)
@@ -1243,7 +1261,10 @@ def pure_gemma_fn(graphdef, state_in, y_in):
         )
         return out_y, nnx.state(merged_layer)
 
-      checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse)
+      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
 
       graphdef, state = nnx.split(self.layers_remainder)
       y, new_state = checkpointed_gemma_fn(graphdef, state, y)

src/maxtext/layers/quantizations.py

@@ -26,7 +26,7 @@
 from aqt.jax.v2 import tiled_dot_general
 from aqt.jax.v2 import calibration
 
-from maxtext.layers import nnx_wrappers
+
 import qwix
 from qwix._src.core import dot_general_qt
 

tests/unit/nnx_decoder_test.py

@@ -533,4 +533,4 @@ def test_different_random_seeds_produce_different_logits(self):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()

tests/unit/tiling_test.py

@@ -236,6 +236,7 @@ def test_vocab_tiling_gradient_with_z_loss(self):
         run_name="grad_test_z_loss_with_tiling",
         enable_checkpointing=False,
         enable_dropout=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=False,
@@ -302,6 +303,7 @@ def test_vocab_tiling_gradient_non_tied_embedding(self):
         run_name="value_and_grad_test_non_tied_with_tiling",
         enable_checkpointing=False,
         enable_dropout=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=False,
@@ -366,6 +368,7 @@ def test_vocab_tiling_gradient_tied_embedding(self):
         self.base_config,
         run_name="grad_test_tied_with_tiling",
         enable_checkpointing=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=True,
@@ -428,6 +431,7 @@ def test_vocab_tiling_gradient_data_parallelism(self):
         run_name="value_and_grad_test_dp_tiling",
         enable_checkpointing=False,
         enable_dropout=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=False,
@@ -492,6 +496,7 @@ def test_vocab_tiling_gradient_tensor_parallelism(self):
         run_name="value_and_grad_test_tp_tiling",
         enable_checkpointing=False,
         enable_dropout=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=False,
@@ -558,6 +563,7 @@ def test_vocab_tiling_gradient_context_parallelism(self):
         run_name="value_and_grad_test_cp_tiling",
         enable_checkpointing=False,
         enable_dropout=False,
+        enable_nnx=False,
         max_target_length=self.seq_len,
         per_device_batch_size=self.batch_size,
         logits_via_embedding=False,