Set NNX flags to true by default

Author: ecnal-cienet

src/maxtext/configs/base.yml

@@ -542,9 +542,13 @@ logical_axis_rules: [
                       ['paged_kv_head_dim_size', []],
                       ['dense_layers', []],
                       ['moe_layers', []],
+                      ['layers_outside_pipeline', []],
+                      ['layers_per_stage', []],
                       ['engram_dim', ['tensor']],
                       ['mhc', []],
                       ['diloco', 'diloco'],
+                      ['num_activations', []],
+                      ['circular_repeats', []],
                     ]
 # Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
 data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']]
@@ -1133,9 +1137,9 @@ position_id_per_seconds: 25
 subslice_shape: ""
 
 # NNX
-enable_nnx: False
-pure_nnx_decoder: False
-pure_nnx: False
+enable_nnx: True
+pure_nnx_decoder: True
+pure_nnx: True
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml

@@ -70,4 +70,32 @@ logical_axis_rules: [
                       ['exp_with_fsdp', 'fsdp'],
                       ['paged_kv_heads', ['tensor']],
                       ['engram_dim', ['tensor']],
+                      # Axes unsharded: sequence/context/tensor_transpose/autoregressive do not exist in this mesh
+                      ['activation_attn_length_no_exp', []],
+                      ['activation_length_no_exp', []],
+                      ['activation_norm_length', []],
+                      ['activation_q_length_no_exp', []],
+                      ['prefill_activation_length', []],
+                      ['prefill_activation_norm_length', []],
+                      ['activation_kv_length', []],
+                      ['decode_length', []],
+                      ['embed_tensor_transpose', []],
+                      ['q_lora_up_proj', []],
+                      ['kv_lora_up_proj', []],
+                      ['kv', []],
+                      ['qkv', []],
+                      ['kv_head_dim', []],
+                      ['cache_batch_prefill', []],
+                      ['cache_batch', []],
+                      ['cache_heads_none', []],
+                      ['cache_kv', []],
+                      ['cache_sequence', []],
+                      ['num_pages', []],
+                      ['tokens_per_page', []],
+                      ['paged_kv_head_dim_size', []],
+                      ['dense_layers', []],
+                      ['moe_layers', []],
+                      ['num_activations', []],
+                      ['mhc', []],
+                      ['diloco', []],
                     ]

src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml

@@ -32,4 +32,57 @@ logical_axis_rules: [
                       ['q_lora', ['fsdp']],
                       ['kv_lora', ['fsdp']],
                       ['exp_with_fsdp', 'fsdp'],
+                      # All other axes are unsharded (tensor/sequence/expert axes do not exist in pure-fsdp)
+                      ['activation_heads', []],
+                      ['activation_kv_heads', []],
+                      ['activation_length', []],
+                      ['activation_attn_length', []],
+                      ['activation_attn_length_no_exp', []],
+                      ['activation_length_no_exp', []],
+                      ['activation_norm_length', []],
+                      ['activation_q_length', []],
+                      ['activation_q_length_no_exp', []],
+                      ['prefill_activation_length', []],
+                      ['prefill_activation_norm_length', []],
+                      ['activation_kv_length', []],
+                      ['activation_attn_embed', []],
+                      ['activation_embed', []],
+                      ['activation_mlp', []],
+                      ['activation_kv', []],
+                      ['activation_kv_head_dim', []],
+                      ['activation_vocab', []],
+                      ['activation_stage', []],
+                      ['activation_exp', []],
+                      ['decode_length', []],
+                      ['mlp', []],
+                      ['mlp_no_fsdp', []],
+                      ['vocab', []],
+                      ['heads', []],
+                      ['q_heads', []],
+                      ['kv_heads', []],
+                      ['embed_tensor_transpose', []],
+                      ['q_lora_up_proj', []],
+                      ['kv_lora_up_proj', []],
+                      ['norm', []],
+                      ['layers', []],
+                      ['qkv', []],
+                      ['kv', []],
+                      ['kv_head_dim', []],
+                      ['cache_batch_prefill', []],
+                      ['cache_batch', []],
+                      ['cache_heads_none', []],
+                      ['cache_heads', []],
+                      ['cache_kv', []],
+                      ['cache_sequence', []],
+                      ['exp', []],
+                      ['paged_kv_heads', []],
+                      ['num_pages', []],
+                      ['tokens_per_page', []],
+                      ['paged_kv_head_dim_size', []],
+                      ['dense_layers', []],
+                      ['moe_layers', []],
+                      ['num_activations', []],
+                      ['engram_dim', []],
+                      ['mhc', []],
+                      ['diloco', []],
                     ]

src/maxtext/configs/decoupled_base_test.yml

@@ -1,6 +1,7 @@
 # Decoupled base test config: used when DECOUPLE_GCLOUD=TRUE for tests that previously relied on base.yml.
-# Inherit all model defaults (PyDantic already does this) but override any cloud-coupled paths and disable 
-# optional cloud features.
+# Inherits from base.yml so that logical_axis_rules, mesh_axes, NNX flags, and all other
+# model defaults are kept in sync. Overrides only cloud-coupled paths and optional cloud features.
+base_config: base.yml
 
 # Output goes to a local relative directory so tests do not require GCS.
 base_output_directory: ./maxtext_local_output/gcloud_decoupled_test_logs
@@ -34,34 +35,9 @@ attention: "dot_product"
 dump_hlo: false
 jax_cache_dir: ""
 
-# Neutral parallelism (single device) for local tests.
-ici_data_parallelism: 1
-ici_tensor_parallelism: 1
-ici_pipeline_parallelism: 1
-ici_expert_parallelism: 1
-ici_sequence_parallelism: 1
-ici_context_parallelism: 1
-ici_tensor_transpose_parallelism: 1
-ici_tensor_sequence_parallelism: 1
-ici_autoregressive_parallelism: 1
-ici_fsdp_parallelism: 1
-ici_fsdp_transpose_parallelism: 1
 # Allow higher unsharded parameter percentage for small device count
 sharding_tolerance: 0.3
 
-# DCN dimensions to 1 (no multi-slice expectation locally).
-dcn_data_parallelism: 1
-dcn_tensor_parallelism: 1
-dcn_pipeline_parallelism: 1
-dcn_expert_parallelism: 1
-dcn_sequence_parallelism: 1
-dcn_context_parallelism: 1
-dcn_tensor_transpose_parallelism: 1
-dcn_tensor_sequence_parallelism: 1
-dcn_autoregressive_parallelism: 1
-dcn_fsdp_parallelism: 1
-dcn_fsdp_transpose_parallelism: 1
-
 # Config logging off unless a test overrides.
 log_config: false
 

src/maxtext/layers/attentions.py

@@ -525,14 +525,14 @@ def __init__(
     elif self.is_qwen3_next:
       self.query_norm = Qwen3NextRMSNorm(
           num_features=self.config.head_dim,
-          eps=self.config.normalization_layer_epsilon,
+          epsilon=self.config.normalization_layer_epsilon,
           dtype=self.config.dtype,
           weight_dtype=self.config.weight_dtype,
           rngs=self.rngs,
       )
       self.key_norm = Qwen3NextRMSNorm(
           num_features=self.config.head_dim,
-          eps=self.config.normalization_layer_epsilon,
+          epsilon=self.config.normalization_layer_epsilon,
           dtype=self.config.dtype,
           weight_dtype=self.config.weight_dtype,
           rngs=self.rngs,

src/maxtext/layers/moe.py

@@ -2041,15 +2041,15 @@ def __call__(
       w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)
       wo_kernel = jnp.asarray(self.wo[...], self.dtype)
 
-    if self.per_expert_scale is not None:
-      wo_kernel = wo_kernel * jnp.asarray(self.per_expert_scale[...], self.dtype)[:, None, None]
+      if self.per_expert_scale is not None:
+        wo_kernel = wo_kernel * jnp.asarray(self.per_expert_scale[...], self.dtype)[:, None, None]
 
-    if cfg.mlp_bias:
-      w0_bias = jnp.asarray(self.wi_0_bias[...], self.dtype)
-      w1_bias = jnp.asarray(self.wi_1_bias[...], self.dtype)
-      wo_bias = jnp.asarray(self.wo_bias[...], self.dtype)
-    else:
-      w0_bias, w1_bias, wo_bias = None, None, None
+      if cfg.mlp_bias:
+        w0_bias = jnp.asarray(self.wi_0_bias[...], self.dtype)
+        w1_bias = jnp.asarray(self.wi_1_bias[...], self.dtype)
+        wo_bias = jnp.asarray(self.wo_bias[...], self.dtype)
+      else:
+        w0_bias, w1_bias, wo_bias = None, None, None
 
       if cfg.sparse_matmul:
         if quantizations.in_serve_mode(self.quant):

src/maxtext/layers/nnx_decoders.py

@@ -303,6 +303,7 @@ def layer_fn(carry, scanned_vars):
       layer = nnx.merge(graphdef, current_params, current_state)
       layer_out = layer(carry, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs)
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+      nnx.pop(layer, nnx.Intermediate)
       return new_carry, nnx.state(layer)
 
     final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state))
@@ -534,6 +535,8 @@ def _create_scanned_layers(
       self, decoder_layer_class, length: int, metadata_axis_name: str, rngs: nnx.Rngs, **layer_kwargs
   ):
     """Creates a VMapped stack of layers, forcing parameter init for Compact modules."""
+    if length == 0:
+      return nnx.List([])
 
     def create_layer_fn(rng):
       return decoder_layer_class(
@@ -566,13 +569,17 @@ 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._uses_linen_fp8_ops():
+      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
 
   def _apply_layers_sequentially(self, layers, x_in, *args, length: int, **kwargs):
     """Runs the layer stack using nnx.scan."""
+    if length == 0:
+      _, empty_state = nnx.split(layers)
+      return x_in, empty_state
     policy = self.get_remat_policy()
     prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
     graphdef, params, state = nnx.split(
@@ -608,7 +615,25 @@ def layer_fn(carry, scanned_vars):
       # Run the layer (Filter kwargs if using the solution from previous turn)
       layer_out = layer(carry, *args, **valid_kwargs)
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
-      return new_carry, nnx.state(layer)
+      nnx.pop(layer, nnx.Intermediate)
+      new_current_state = nnx.state(layer)
+      return new_carry, new_current_state
+
+    if self._uses_linen_fp8_ops():
+      # jax.lax.scan is incompatible with Linen fp8 ops: put_variable in setup() stores
+      # scan-level tracers as Python attributes on the Linen module, causing a tracer leak
+      # across the scan boundary. Fall back to a Python loop instead.
+      x = x_in
+      for i in range(length):
+        params_i = jax.tree.map(lambda p, _i=i: p[_i], params)
+        state_i = jax.tree.map(lambda s, _i=i: s[_i], state)
+        layer = nnx.merge(graphdef, params_i, state_i)
+        layer_out = layer(x, *args, **valid_kwargs)
+        x = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+        nnx.pop(layer, nnx.Intermediate)
+      if scan_axis != 0:
+        params = jax.tree.map(lambda p: jnp.moveaxis(p, 0, scan_axis), params)
+      return x, nnx.State.merge(params, state)
 
     layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
     final_carry, scanned_state = jax.lax.scan(layer_fn, x_in, (params, state))
@@ -672,7 +697,8 @@ def get_chunk(pytree, start, end):
         layer_out = layer(y, *layer_args, **valid_kwargs)
         y = layer_out[0] if isinstance(layer_out, tuple) else layer_out
 
-        _, new_eng_mutables = nnx.split(layer, nnx.Param, ...)
+        nnx.pop(layer, nnx.Intermediate)
+        _, _, new_eng_mutables = nnx.split(layer, nnx.Param, ...)
         new_eng_mutables = jax.tree.map(lambda x: jnp.expand_dims(x, axis=0), new_eng_mutables)
         updated_mutables_chunks.append(new_eng_mutables)
         current_idx += 1
@@ -698,10 +724,12 @@ def layer_fn(carry, scanned_vars):
           l = nnx.merge(graphdef, curr_p, curr_m)
           l_out = l(carry, *layer_args, **valid_kwargs)
           n_carry = l_out[0] if isinstance(l_out, tuple) else l_out
-          _, n_mut = nnx.split(l, nnx.Param, ...)
+          nnx.pop(l, nnx.Intermediate)
+          _, _, n_mut = nnx.split(l, nnx.Param, ...)
           return n_carry, n_mut
 
-        layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+        if not self._uses_linen_fp8_ops():
+          layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
         y, new_chunk_mutables = jax.lax.scan(layer_fn, y, (chunk_params, chunk_mutables))
         updated_mutables_chunks.append(new_chunk_mutables)
         current_idx = next_boundary
@@ -742,7 +770,11 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in, dynamic_kwargs):
       out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **dynamic_kwargs)
       return out_y, out_kv, nnx.state(merged_layer)
 
-    checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
+    checkpointed_fn = (
+        pure_layer_fn
+        if self._uses_linen_fp8_ops()
+        else jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
+    )
 
     for lyr in range(num_layers):
       attr_name = f"{base_name}_{lyr}"
@@ -921,6 +953,10 @@ def get_remat_policy(self):
         assert cfg.remat_policy == "full", "Remat policy needs to be on list of remat policies"
     return policy
 
+  def _uses_linen_fp8_ops(self) -> bool:
+    """Returns True if the quantization mode uses Linen fp8 ops incompatible with jax.checkpoint."""
+    return self.config.quantization in ("fp8_gpu", "fp8_nanoo")
+
   def get_norm_layer(self, num_features: int, rngs: nnx.Rngs):
     """Helper to retrieve the correct normalization layer class based on config, partially applied with common arguments."""
     if self.config.decoder_block in (
@@ -1072,10 +1108,18 @@ def __call__(
       audio_embeddings: None | jnp.ndarray = None,
       audio_masks: None | jnp.ndarray = None,
       deepstack_visual_embeds: None | list[jnp.ndarray] = None,
+      multimodal_input=None,
   ):
     cfg = self.config
     assert decoder_input_tokens.ndim == 2  # [batch, len]
 
+    if multimodal_input is not None:
+      image_embeddings = multimodal_input.image_embeddings
+      bidirectional_mask = multimodal_input.bidirectional_mask
+      image_masks = multimodal_input.image_masks
+      audio_embeddings = multimodal_input.audio_embeddings
+      audio_masks = multimodal_input.audio_masks
+
     # [batch, length] -> [batch, length, emb_dim]
     y = self._apply_embedding(
         shared_embedding,
@@ -1223,7 +1267,6 @@ def __call__(
                   decoder_input_tokens=decoder_input_tokens,
               )
 
-
     else:
       # Non-Pipeline Run
       if cfg.scan_layers:

src/maxtext/layers/nnx_wrappers.py

@@ -26,6 +26,7 @@
 from flax.core import FrozenDict
 from flax.core import meta
 from flax.nnx import graph
+from flax.nnx import tracers as nnx_tracers
 from flax.nnx import variablelib
 from flax.nnx.bridge import module as bdg_module
 from flax.nnx.module import Module
@@ -170,6 +171,23 @@ def current_linen_module() -> linen.Module | None:
   return None
 
 
+def _refresh_variable_trace_state(module: Module) -> None:
+  """Refresh _trace_state for Variables that have stale trace state.
+
+  When nnx.update() is called with tracer values from a JAX transformation
+  (e.g. jax.grad's LinearizeTracer), it uses _unsafe_bypass_check=True which
+  updates the raw value but not _trace_state. This leaves Variables with a
+  stale _trace_state from the outer (Python) context, causing nnx.split() to
+  fail with "Cannot extract graph node from different trace level" errors.
+
+  This function resets _trace_state on any Variables whose _can_update is False
+  so that downstream NNX operations (e.g. nnx.split in NNXPipeline) succeed.
+  """
+  for _, v in nnx.graph.iter_graph(module):
+    if isinstance(v, variablelib.Variable) and not v._can_update:  # pylint: disable=protected-access
+      object.__setattr__(v, "_trace_state", nnx_tracers.TraceState())
+
+
 class ToNNX(Module):
   """A wrapper to turn any Linen module into an NNX module.
 
@@ -467,6 +485,7 @@ def maybe_unbox(x):
       warnings.warn(f"Found unknown module paths in incoming state:{paths_str}")
 
     nnx.update(module, new_state)
+    _refresh_variable_trace_state(module)
 
     _fix_for_qwix_quantization(module)
     method_fn = _get_module_method(module, nnx_method)