Merge pull request #4507 from AI-Hypercomputer:aireen/pr-nnx-scan-helper

PiperOrigin-RevId: 949755622

Author: Google-ML-Automation

src/maxtext/layers/nnx_decoders.py

@@ -37,7 +37,7 @@
     ShardMode,
 )
 from maxtext.layers import initializers, linears, mhc, normalizations, quantizations
-from maxtext.layers import nnx_wrappers
+from maxtext.layers import nnx_scan, nnx_wrappers
 from maxtext.layers.attentions import Attention
 from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding
 from maxtext.layers.normalizations import RMSNorm
@@ -839,94 +839,20 @@ def _create_scanned_layers(
       **layer_kwargs,
   ):
     """Creates a scanned stack of layers using jax.lax.scan for memory-efficient initialization."""
-    if length == 0:
-      return None
-    scan_axis = self.config.param_scan_axis
-
-    # Fork rngs to get per-layer RNG states for scanning
-    try:
-      forked_rngs = rngs.fork(split=length)
-    except:  # pylint: disable=bare-except
-      pass
-
-    rngs_graphdef, rngs_state = nnx.split(forked_rngs)
-
-    first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
-    ref_rngs = nnx.merge(rngs_graphdef, first_rng_state)
-    ref_layer = decoder_layer_class(
-        config=self.config,
-        mesh=self.mesh,
-        quant=self.quant,
-        model_mode=self.model_mode,
-        rngs=ref_rngs,
-        **layer_kwargs,
+    return nnx_scan.create_scanned_layers(
+        lambda layer_rngs: decoder_layer_class(
+            config=self.config,
+            mesh=self.mesh,
+            model_mode=self.model_mode,
+            quant=self.quant,
+            rngs=layer_rngs,
+            **layer_kwargs,
+        ),
+        length=length,
+        param_scan_axis=self.config.param_scan_axis,
+        metadata_axis_name=metadata_axis_name,
+        rngs=rngs,
     )
-    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)
-
-    _, (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)
-
-    def _add_scan_metadata(state, axis):
-      def _update_leaf(leaf):
-        if hasattr(leaf, "replace") and hasattr(leaf, "value"):
-          replace_kwargs = {}
-          if hasattr(leaf, "get_metadata"):
-            replace_kwargs.update(leaf.get_metadata())
-
-          replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name
-          replace_kwargs["param_scan_axis"] = axis
-
-          for key in [
-              "sharding",
-              "out_sharding",
-              "kernel_axes",
-              "sharding_names",
-          ]:
-            val = getattr(leaf, key, None)
-            if val is None and key in replace_kwargs:
-              val = replace_kwargs[key]
-
-            if val is not None:
-              if isinstance(val, str):
-                val = (val,)
-              if isinstance(val, tuple):
-                l = list(val)
-                # Safely insert the scan axis into the logical axes string
-                if metadata_axis_name not in l:
-                  insert_idx = min(axis, len(l))
-                  l.insert(insert_idx, metadata_axis_name)
-                  replace_kwargs[key] = tuple(l)
-
-          return leaf.replace(**replace_kwargs)
-        return leaf
-
-      # We must use a custom is_leaf to catch the VariableState instances
-      return jax.tree.map(
-          _update_leaf,
-          state,
-          is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),
-      )
-
-    stacked_params = _add_scan_metadata(stacked_params, scan_axis)
-    stacked_rest = _add_scan_metadata(stacked_rest, 0)
-
-    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."""

src/maxtext/layers/nnx_scan.py

@@ -0,0 +1,89 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for constructing stacks of scanned NNX layers."""
+
+from collections.abc import Callable
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+
+
+def create_scanned_layers(
+    layer_factory: Callable[[nnx.Rngs], nnx.Module],
+    *,
+    length: int,
+    param_scan_axis: int,
+    metadata_axis_name: str,
+    rngs: nnx.Rngs,
+) -> nnx.Module | None:
+  """Constructs an NNX layer whose variables are stacked for a layer scan."""
+  if length == 0:
+    return None
+
+  forked_rngs = rngs.fork(split=length)
+  rngs_graphdef, rngs_state = nnx.split(forked_rngs)
+
+  first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
+  reference_layer = layer_factory(nnx.merge(rngs_graphdef, first_rng_state))
+  layer_graphdef, _, _ = nnx.split(reference_layer, nnx.Param, ...)
+  del reference_layer
+
+  def scan_body(carry, rng_state_slice):
+    layer = layer_factory(nnx.merge(rngs_graphdef, rng_state_slice))
+    _, params, rest = nnx.split(layer, nnx.Param, ...)
+    return carry, (params, rest)
+
+  _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
+
+  if param_scan_axis != 0:
+    stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), stacked_params)
+
+  def add_scan_metadata(state, axis):
+    def update_leaf(leaf):
+      if hasattr(leaf, "replace") and hasattr(leaf, "value"):
+        replace_kwargs = {}
+        if hasattr(leaf, "get_metadata"):
+          replace_kwargs.update(leaf.get_metadata())
+
+        replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name
+        replace_kwargs["param_scan_axis"] = axis
+
+        for key in ["sharding", "out_sharding", "kernel_axes", "sharding_names"]:
+          value = getattr(leaf, key, None)
+          if value is None and key in replace_kwargs:
+            value = replace_kwargs[key]
+
+          if value is not None:
+            if isinstance(value, str):
+              value = (value,)
+            if isinstance(value, tuple):
+              logical_axes = list(value)
+              if metadata_axis_name not in logical_axes:
+                logical_axes.insert(min(axis, len(logical_axes)), metadata_axis_name)
+                replace_kwargs[key] = tuple(logical_axes)
+
+        return leaf.replace(**replace_kwargs)
+      return leaf
+
+    return jax.tree.map(
+        update_leaf,
+        state,
+        is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),
+    )
+
+  stacked_params = add_scan_metadata(stacked_params, param_scan_axis)
+  stacked_rest = add_scan_metadata(stacked_rest, 0)
+  return nnx.merge(layer_graphdef, stacked_params, stacked_rest)

tests/unit/nnx_scan_test.py

@@ -0,0 +1,65 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for NNX scan utilities."""
+
+import unittest
+
+from flax import nnx
+import jax
+import pytest
+
+from maxtext.layers import nnx_scan
+
+
+class _LinearLayer(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs):
+    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (2, 2)))
+
+  def __call__(self, inputs):
+    return inputs @ self.kernel.value
+
+
+@pytest.mark.cpu_only
+class TestCreateScannedLayers(unittest.TestCase):
+  """Tests for nnx_scan.create_scanned_layers."""
+
+  def test_create_stacks_params_at_param_scan_axis(self):
+    """Per-layer params are stacked along param_scan_axis."""
+    length = 3
+    for axis, expected_shape in ((0, (length, 2, 2)), (1, (2, length, 2))):
+      layers = nnx_scan.create_scanned_layers(
+          _LinearLayer,
+          length=length,
+          param_scan_axis=axis,
+          metadata_axis_name="layers",
+          rngs=nnx.Rngs(0),
+      )
+      self.assertEqual(layers.kernel.value.shape, expected_shape)
+
+  def test_create_zero_length_returns_none(self):
+    """A zero-length stack short-circuits to None."""
+    layers = nnx_scan.create_scanned_layers(
+        _LinearLayer,
+        length=0,
+        param_scan_axis=0,
+        metadata_axis_name="layers",
+        rngs=nnx.Rngs(0),
+    )
+    self.assertIsNone(layers)
+
+
+if __name__ == "__main__":
+  unittest.main()