[WIP] NNX: fix model and test compatibility issues

- Replace nn.Dropout with linears.Dropout in gpt_oss and olmo3 decoder layers
- Add num_activations logical axis rule to base.yml
- Fix integration and unit tests for NNX compatibility

I will relocate these files accordingly once the work is done.

Author: ecnal-cienet

src/maxtext/common/gcloud_stub.py

@@ -43,6 +43,15 @@ def is_decoupled() -> bool:  # dynamic check so setting env after initial import
   return os.environ.get("DECOUPLE_GCLOUD", "").upper() == "TRUE"
 
 
+def is_pure_nnx() -> bool:  # dynamic check so setting env after initial import still works
+  """Return True when running in pure NNX mode (PURE_NNX=TRUE env var).
+
+  Defaults to FALSE — Linen is the default test mode.
+  Set PURE_NNX=TRUE to opt in to NNX mode (skips linen_only tests, runs nnx_only tests).
+  """
+  return os.environ.get("PURE_NNX", "FALSE").upper() == "TRUE"
+
+
 T = TypeVar("T")
 
 

src/maxtext/configs/base.yml

@@ -514,6 +514,7 @@ logical_axis_rules: [
                       ['paged_kv_head_dim_size', []],
                       ['dense_layers', []],
                       ['moe_layers', []],
+                      ['num_activations', []],
                       ['engram_dim', ['tensor']],
                       ['mhc', []],
                       ['diloco', 'diloco'],

src/maxtext/configs/decoupled_base_test.yml

@@ -30,6 +30,10 @@ eval_dataset_name: 'c4/en:3.1.0'
 # Use dot_product attention to avoid GPU Pallas shared memory limits on AMD GPUs
 attention: "dot_product"
 
+# Default to Linen mode for tests; NNX is opt-in via PURE_NNX=TRUE.
+pure_nnx: False
+pure_nnx_decoder: False
+
 # Avoid HLO dump overhead.
 dump_hlo: false
 jax_cache_dir: ""

src/maxtext/layers/nnx_decoders.py

@@ -432,8 +432,16 @@ 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)
+    # Linen-based FP8 ops (fp8_nanoo, fp8_gpu) store scale/amax_history in Linen
+    # mutable scope. jax.checkpoint re-traces the scan body during backward (remat),
+    # but the Linen scope retains JAX tracers from the first trace, causing
+    # UnexpectedTracerError. Skip checkpoint for these quantization types.
+    uses_linen_fp8_mutable_state = self.config.quantization in ("fp8_nanoo", "fp8_gpu")
+    if uses_linen_fp8_mutable_state:
+      out, new_state = pure_layer_fn(state, y)
+    else:
+      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
@@ -475,20 +483,29 @@ def layer_fn(carry, scanned_vars):
 
       return new_carry, new_current_state
 
-    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))
+    # Linen-based FP8 ops (fp8_nanoo, fp8_gpu) store scale/amax_history in Linen
+    # mutable scope. jax.lax.scan traces the body function and Linen's setup() creates
+    # intermediate tracer values (amax_history float32[1024]) that escape the scan scope,
+    # causing UnexpectedTracerError. Use a Python for loop instead for these types.
+    uses_linen_fp8_mutable_state = self.config.quantization in ("fp8_nanoo", "fp8_gpu")
+    if uses_linen_fp8_mutable_state:
+      carry = x_in
+      per_layer_states = []
+      for i in range(length):
+        current_params = jax.tree.map(lambda x, i=i: x[i], params)
+        current_state = jax.tree.map(lambda x, i=i: x[i], state)
+        carry, new_state_i = layer_fn(carry, (current_params, current_state))
+        per_layer_states.append(new_state_i)
+      final_carry = carry
+      scanned_state = jax.tree.map(lambda *xs: jnp.stack(list(xs)), *per_layer_states)
+    else:
+      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))
 
     if scan_axis != 0:
       params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), params)
 
-    # Params are read-only during the forward pass, so the scan output's copy of
-    # params is at axis=0 (lax.scan default) rather than scan_axis. Discard the
-    # scan-output params and keep the original params (correctly positioned at
-    # scan_axis) to avoid a shape mismatch when _apply_scanned_chunk tries to
-    # write them back via dynamic_update_slice_in_dim.
-    _, non_param_scanned_state = scanned_state.split(nnx.Param, ...)
-    scanned_state = nnx.State.merge(params, non_param_scanned_state)
+    scanned_state = nnx.State.merge(params, scanned_state)
     return final_carry, nnx.merge(graphdef, scanned_state)
 
   def get_decoder_layers(self):

src/maxtext/layers/normalizations.py

@@ -104,9 +104,9 @@ def __call__(self, x: jnp.ndarray, out_sharding: NamedSharding | None = None) ->
 
 def Qwen3NextRMSNorm(
     num_features: int,
-    epsilon: float,
-    dtype: DType,
-    weight_dtype: DType,
+    epsilon: float = 1e-6,
+    dtype: DType = jnp.float32,
+    weight_dtype: DType = jnp.float32,
     shard_mode: ShardMode = ShardMode.AUTO,
     kernel_axes: tuple[None | str, ...] = (),
     parameter_memory_host_offload: bool = False,

src/maxtext/models/gpt_oss.py

@@ -28,6 +28,7 @@
 from maxtext.common.common_types import AttentionType, Config
 from maxtext.layers import attentions
 from maxtext.layers import initializers
+from maxtext.layers import linears
 from maxtext.layers import moe
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import quantizations
@@ -130,6 +131,8 @@ def __init__(
         rngs=rngs,
     )
 
+    self.dropout = linears.Dropout(rate=config.dropout_rate, broadcast_dims=(-2,), rngs=rngs)
+
   def __call__(
       self,
       inputs,
@@ -181,7 +184,7 @@ def __call__(
     mlp_lnx = nn.with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed"))
 
     layer_output = mlp_lnx + intermediate_inputs
-    layer_output = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(layer_output, deterministic=deterministic)
+    layer_output = self.dropout(layer_output, deterministic=deterministic)
 
     layer_output = nn.with_logical_constraint(
         layer_output,

src/maxtext/models/llama2.py

@@ -70,6 +70,7 @@ def __init__(
         shard_mode=config.shard_mode,
         kernel_axes=("norm",),
         epsilon=config.normalization_layer_epsilon,
+        parameter_memory_host_offload=config.parameter_memory_host_offload,
         rngs=rngs,
     )
 

src/maxtext/models/olmo3.py

@@ -29,6 +29,7 @@
 from maxtext.common.common_types import AttentionType, Config
 from maxtext.layers import attentions
 from maxtext.layers import initializers
+from maxtext.layers import linears
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import quantizations
 from maxtext.layers.attentions import Attention
@@ -140,6 +141,8 @@ def __init__(
         rngs=rngs,
     )
 
+    self.dropout = linears.Dropout(rate=config.dropout_rate, broadcast_dims=(-2,), rngs=rngs)
+
   def __call__(
       self,
       inputs,
@@ -193,7 +196,7 @@ def __call__(
     mlp_lnx = nn.with_logical_constraint(mlp_lnx, ("activation_batch", "activation_norm_length", "activation_embed"))
 
     layer_output = mlp_lnx + intermediate_inputs
-    layer_output = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(layer_output, deterministic=deterministic)
+    layer_output = self.dropout(layer_output, deterministic=deterministic)
 
     layer_output = nn.with_logical_constraint(
         layer_output,

src/maxtext/trainers/post_train/sft/train_sft.py

@@ -47,8 +47,6 @@
 
 from orbax import checkpoint as ocp
 
-from tunix.sft import metrics_logger, peft_trainer, profiler
-
 from maxtext.configs import pyconfig
 from maxtext.trainers.pre_train.train import loss_fn
 from maxtext.common.goodput import (
@@ -77,6 +75,8 @@ def get_tunix_config(mt_config):
   Returns:
     A Tunix `TrainingConfig` object.
   """
+  from tunix.sft import metrics_logger, peft_trainer, profiler  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+
   # Checkpointing configurations
   checkpointing_options = ocp.CheckpointManagerOptions(
       save_interval_steps=mt_config.checkpoint_period,
@@ -143,6 +143,8 @@ def loss_func(model, inputs, inputs_position, inputs_segmentation, targets, targ
 
 def setup_trainer_state(mt_config, goodput_recorder=None):
   """Set up prerequisites for training loop."""
+  from tunix.sft import peft_trainer  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+
   tunix_config = get_tunix_config(mt_config)
 
   with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT):

src/maxtext/trainers/pre_train/train.py

@@ -497,7 +497,11 @@ def move(path, value):
     if config.use_dpo:
       new_state = _merge_dpo_state(new_state, reference_params)
     return new_state, metrics
-  return nnx.state(new_state), metrics
+  # Exclude Intermediate variables (e.g., sowed max_logits for QK-Clip) from the
+  # returned state. Intermediates are transient forward-pass artifacts and must not
+  # persist across steps: they're absent from the abstract state used to build
+  # state_mesh_shardings, so including them would cause a leaf-count mismatch in JAX.
+  return nnx.state(new_state, nnx.Not(nnx.Intermediate)), metrics
 
 
 def eval_step(model, config, state, data, dropout_rng=None):