fix(nnx): add eager sharding and python loop initialization to resolve compilation OOM

This commit introduces:
1. An eager sharding monkey-patch for nnx.Param to avoid replicating weights on host during initialization.
2. Python loop initialization for decoder layers (instead of jax.lax.scan) when eager sharding is enabled, to prevent compilation memory spike.
3. A configuration flag `enable_nnx_scan_oom_fix` to control these optimizations.

TAG=agy
CONV=0c977f36-644a-472a-981c-1856f8cac8a6

Author: hsuan-lun-chiang

src/maxtext/configs/base.yml

@@ -1236,6 +1236,7 @@ subslice_shape: ""
 enable_nnx: true
 pure_nnx_decoder: true
 pure_nnx: true
+enable_nnx_scan_oom_fix: true
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/maxtext/configs/types.py

@@ -1017,6 +1017,7 @@ class HardwareAndMesh(BaseModel):
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
   pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.")
   pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.")
+  enable_nnx_scan_oom_fix: bool = Field(True, description="Whether to enable the OOM fix for NNX scanned decoder.")
   remove_size_one_mesh_axis_from_type: bool = Field(
       True,
       description="Whether to remove size one mesh axis from type through jax.config.",

src/maxtext/layers/nnx_decoders.py

@@ -358,8 +358,11 @@ def layer_fn(carry, scanned_vars):
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
       # Avoid returning and stacking read-only parameters inside the scan body.
       # This prevents huge unnecessary memory allocation.
-      _, _, updated_state = nnx.split(layer, nnx.Param, ...)
-      return new_carry, updated_state
+      if self.config.enable_nnx_scan_oom_fix:
+        _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+        return new_carry, updated_state
+      else:
+        return new_carry, nnx.state(layer)
 
     final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state))
 
@@ -868,20 +871,40 @@ def _create_scanned_layers(
     layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...)
     del ref_layer
 
-    def scan_body(carry, rng_state_slice):
-      layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice)
-      layer = decoder_layer_class(
-          config=self.config,
-          mesh=self.mesh,
-          quant=self.quant,
-          model_mode=self.model_mode,
-          rngs=layer_rngs,
-          **layer_kwargs,
-      )
-      _, params, rest = nnx.split(layer, nnx.Param, ...)
-      return carry, (params, rest)
+    if self.config.enable_nnx_scan_oom_fix:
+      params_list = []
+      rest_list = []
+      for i in range(length):
+        layer_rng_state = jax.tree.map(lambda x: x[i], rngs_state)
+        layer_rngs = nnx.merge(rngs_graphdef, layer_rng_state)
+        layer = decoder_layer_class(
+            config=self.config,
+            mesh=self.mesh,
+            quant=self.quant,
+            model_mode=self.model_mode,
+            rngs=layer_rngs,
+            **layer_kwargs,
+        )
+        _, params, rest = nnx.split(layer, nnx.Param, ...)
+        params_list.append(params)
+        rest_list.append(rest)
+      stacked_params = jax.tree.map(lambda *args: jnp.stack(args), *params_list)
+      stacked_rest = jax.tree.map(lambda *args: jnp.stack(args), *rest_list)
+    else:
+      def scan_body(carry, rng_state_slice):
+        layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice)
+        layer = decoder_layer_class(
+            config=self.config,
+            mesh=self.mesh,
+            quant=self.quant,
+            model_mode=self.model_mode,
+            rngs=layer_rngs,
+            **layer_kwargs,
+        )
+        _, params, rest = nnx.split(layer, nnx.Param, ...)
+        return carry, (params, rest)
 
-    _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
+      _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
 
     if scan_axis != 0:
       stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), stacked_params)
@@ -1037,8 +1060,11 @@ def layer_fn(carry, scanned_vars):
       else:
         # Avoid returning and stacking read-only parameters inside the scan body.
         # This prevents huge unnecessary memory allocation.
-        _, _, updated_state = nnx.split(layer, nnx.Param, ...)
-        new_current_state = updated_state
+        if self.config.enable_nnx_scan_oom_fix:
+          _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+          new_current_state = updated_state
+        else:
+          new_current_state = nnx.state(layer)
 
       if use_kv:
         return new_carry, (new_current_state, updated_kv)

src/maxtext/utils/model_creation_utils.py

@@ -523,6 +523,9 @@ def from_config(
   """
   if mesh is None:
     mesh = maxtext_utils.get_mesh_from_config(config, devices)
+  from maxtext.utils import sharding
+  sharding.active_config = config
+  sharding.active_mesh = mesh
   model = create_model(config, mesh, model_mode=model_mode, rngs=rngs, quant_mode_str=quant_mode_str)
 
   # Return only the model

src/maxtext/utils/sharding.py

@@ -919,3 +919,25 @@ def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, sha
   # Apply the constraint to the model's current variables. This tells JAX to
   # gather the weights into this layout.
   return maybe_shard_with_name(variables, physical_constraint_no_fsdp, shard_mode=shard_mode)
+
+
+# Global references for monkey-patched nnx.Param initialization
+active_config = None
+active_mesh = None
+
+original_param_init = nnx.Param.__init__
+
+def patched_param_init(self, value, *args, **kwargs):
+  sharding_axes = kwargs.get('sharding') or kwargs.get('out_sharding')
+  if sharding_axes is not None:
+    if active_config is not None and active_mesh is not None:
+      value = maybe_shard_with_logical(
+          value,
+          sharding_axes,
+          active_mesh,
+          active_config.shard_mode,
+          rules=active_config.logical_axis_rules,
+      )
+  original_param_init(self, value, *args, **kwargs)
+
+nnx.Param.__init__ = patched_param_init

tests/utils/forward_pass_logit_checker.py

@@ -233,12 +233,12 @@ def check_kl_divergence(model_logits, golden_logits, atol=0.02, clip_logits_epsi
 def get_data(golden_data_point, config):
   """Get the golden data for the test indexed at golden_data_index"""
 
+  model_prefix = config.model_name.split("-")[0]
   max_logging.log(f"config.global_batch_size_to_train_on={config.global_batch_size_to_train_on}")
   if config.use_multimodal:
     assert "pixel_values" in golden_data_point, "no image found in golden data while use_multimodal=True"
     pixel_values = np.asarray(golden_data_point["pixel_values"], dtype=np.float32)
     max_logging.log(f"pixel_values.shape = {pixel_values.shape}")
-    model_prefix = config.model_name.split("-")[0]
     # Gemma3 and Gemma4 models expect (num_images, height, width, channels)
     if model_prefix in ["gemma3", "gemma4"]:
       if pixel_values.ndim == 2: