test mesa

Author: mesakhcienet

src/maxtext/layers/nnx_decoders.py

@@ -394,33 +394,61 @@ def _create_single_layer(self, decoder_layer_class, rngs, **kwargs):
       return nnx_wrappers.ToNNX(layer_linen, rngs=rngs)
 
   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."""
+    """Creates a scanned stack of layers using jax.lax.scan for memory-efficient sequential initialization.
 
-    def create_layer_fn(rng):
+    Uses jax.lax.scan instead of nnx.vmap to reduce peak memory during initialization.
+    With vmap, all layers' parameters are created simultaneously (O(N) peak memory).
+    With scan, parameters are created one layer at a time (O(1) peak intermediate memory),
+    which prevents OOM on memory-constrained devices like TPU v6e-4.
+    """
+    scan_axis = self.config.param_scan_axis
+
+    # Split rngs to get per-layer RNG states
+    split_rngs = rngs.split(length)
+    rngs_graphdef, rngs_state = nnx.split(split_rngs)
+
+    # Create a reference layer to capture the module graph structure (graphdef).
+    # This layer's params are discarded — only the structure is kept.
+    ref_rngs = nnx.Rngs(0)
+    ref_layer = decoder_layer_class(
+        config=self.config, mesh=self.mesh, quant=self.quant,
+        model_mode=self.model_mode, rngs=ref_rngs, **layer_kwargs
+    )
+    layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...)
+
+    # Sequentially create each layer's parameters via jax.lax.scan.
+    # The scan body is traced once; XLA executes it N times with different RNG keys,
+    # keeping only one layer's intermediate state alive at a time.
+    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=rng, **layer_kwargs
+          config=self.config, mesh=self.mesh, quant=self.quant,
+          model_mode=self.model_mode, rngs=layer_rngs, **layer_kwargs
       )
-      return nnx.split(layer, nnx.Param, ...)
-      # return layer
-
-    try:
-      forked_rngs = rngs.fork(split=length)
-    except:  # pylint: disable=bare-except
-      pass
-
-    graphdef, params, rest = nnx.vmap(
-        create_layer_fn,
-        in_axes=0,
-        out_axes=(None, self.config.param_scan_axis, 0),
-        axis_name=metadata_axis_name,
-        transform_metadata={
-            nnx.PARTITION_NAME: metadata_axis_name,
-            "param_scan_axis": self.config.param_scan_axis,
-        },
-    )(forked_rngs)
-    layers_vmapped = nnx.merge(graphdef, params, rest)
+      _, params, rest = nnx.split(layer, nnx.Param, ...)
+      return carry, (params, rest)
 
-    return layers_vmapped
+    _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
+
+    # jax.lax.scan stacks outputs along axis 0. Move params to the configured scan axis.
+    if scan_axis != 0:
+      stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params)
+
+    # Add partition metadata that nnx.vmap's transform_metadata would normally set.
+    # This metadata is read by variable_to_logically_partitioned() in initializers.py
+    # to insert the scan axis name into logical sharding specs.
+    def _add_partition_metadata(state):
+      def _update(vs):
+        if isinstance(vs, nnx.Variable):
+          metadata = vs.get_metadata()
+          return type(vs)(vs.get_value(), **{**metadata, nnx.PARTITION_NAME: metadata_axis_name, "param_scan_axis": scan_axis})
+        return vs
+      return jax.tree.map(_update, state, is_leaf=lambda x: isinstance(x, nnx.Variable))
+
+    stacked_params = _add_partition_metadata(stacked_params)
+    stacked_rest = _add_partition_metadata(stacked_rest)
+
+    return nnx.merge(layer_graphdef, stacked_params, stacked_rest)
 
   def _apply_layer_with_remat(self, layer: nnx.Module, y: jax.Array, policy: Any, prevent_cse: bool, **kwargs):
     """Helper to cleanly apply jax.checkpoint to a single unscanned layer or block."""