fix(nnx): Fix Linen checkpoint restoration into NNX Gemma4 model

.dockerignore

@@ -1,2 +1,3 @@
 .git
 maxtext_venv
+.venv

src/dependencies/scripts/docker_upload_runner.sh

@@ -109,14 +109,9 @@ if [ -n "$ABSOLUTE_LINKS" ]; then
   exit 1
 fi
 
-# Check if the base image exists locally
-if ! docker image inspect "${LOCAL_IMAGE_NAME}" &> /dev/null; then
-  echo "ERROR: Base image '${LOCAL_IMAGE_NAME}' not found locally."
-  echo "Please build it first by running 'build_maxtext_docker_image'."
-  exit 1
-fi
+BASEIMAGE="${BASEIMAGE:-${LOCAL_IMAGE_NAME}}"
 
-docker build --no-cache --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} \
+docker build --no-cache --build-arg BASEIMAGE=${BASEIMAGE} \
              --build-arg PACKAGE_DIR=${PACKAGE_DIR} \
              -f "$PACKAGE_DIR"'/dependencies/dockerfiles/maxtext_runner.Dockerfile' \
              -t ${LOCAL_IMAGE_NAME_RUNNER} .

src/maxtext/utils/maxtext_utils_nnx.py

@@ -65,13 +65,18 @@ def nnx_extract_named_sharding(abstract_state: nnx.State) -> nnx.State:
     out_shardings). Contrast with sharding.nnx_construct_named_sharding, which
     retains wrappers for abstract tree zipping compatibility.
   """
-  # Don't use nnx.get_named_sharding() because it constructs new shardings. Instead, we
-  # get the existing sharding from the abstract_state.
-  # The state leaf is of type jax.ShapeDtypeStruct(shape, dtype, sharding)
+  def _get_sharding(x):
+    if isinstance(x, nnx.Variable):
+      val = x.get_value() if hasattr(x, "get_value") else x.value
+      return getattr(val, "sharding", None)
+    if isinstance(x, jax.ShapeDtypeStruct):
+      return x.sharding
+    return getattr(x, "sharding", None)
+
   return jax.tree.map(
-      lambda x: x.sharding,
+      _get_sharding,
       abstract_state,
-      is_leaf=lambda x: isinstance(x, jax.ShapeDtypeStruct),
+      is_leaf=lambda x: isinstance(x, (nnx.Variable, jax.ShapeDtypeStruct)),
   )
 
 
@@ -170,13 +175,22 @@ def create_nnx_sharded_model(
   # By providing out_shardings, we instruct JAX to produce sharded output directly,
   # avoiding a large intermediate allocation on a single device.
   @partial(jax.jit, out_shardings=named_sharding)
-  def create_sharded_state():
-    model = init_fn()
-    return jax.lax.with_sharding_constraint(nnx.state(model), named_sharding)
+  def create_sharded_zeros():
+    return jax.tree.map(
+        lambda x: jnp.zeros(x.shape, dtype=x.dtype),
+        abstract_state,
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
 
   # Create the model with sharded parameters.
   with jax.set_mesh(mesh):
-    sharded_state = create_sharded_state()
+    sharded_zeros = create_sharded_zeros()
+    sharded_state = jax.tree.map(
+        lambda var, val: var.replace(value=val),
+        abstract_state,
+        sharded_zeros,
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
   return nnx.merge(graphdef, sharded_state)
 
 

src/maxtext/utils/model_creation_utils.py

@@ -402,6 +402,15 @@ def _lookup_stored_meta(path):
       if raw in node:
         node = node[raw]
         continue
+      if name.startswith("layers_") and name[7:].isdigit() and "layers" in node:
+        idx_str = name[7:]
+        layers_node = node["layers"]
+        if isinstance(layers_node, dict) and idx_str in layers_node:
+          node = layers_node[idx_str]
+          continue
+        elif isinstance(layers_node, (list, tuple)) and int(idx_str) < len(layers_node):
+          node = layers_node[int(idx_str)]
+          continue
       return None
     return node
 
@@ -624,18 +633,30 @@ def create_nnx_sharded_model_hybrid(config, mesh=None, devices=None, model_mode=
   # avoiding a large intermediate allocation on a single device.
   with nn.logical_axis_rules(config.logical_axis_rules):
     out_shardings = nn.logical_to_mesh_sharding(specs, mesh)
+    out_shardings = jax.tree.map(
+        lambda x: x.get_value() if isinstance(x, nnx.Variable) else (x.value if hasattr(x, "value") else x),
+        out_shardings,
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
 
   @partial(jax.jit, out_shardings=out_shardings)
-  def create_sharded_state():
-    # This will be JIT-compiled. JAX knows the output sharding and can
-    # initialize the parameters directly on the target devices in a sharded way.
-    model = _create_model_partial()
-    return nnx.state(model)
+  def create_sharded_zeros():
+    return jax.tree.map(
+        lambda x: jnp.zeros(x.shape, dtype=x.dtype),
+        abstract_state,
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
 
   with mesh:
     # Create the model with sharded parameters.
     with nn.logical_axis_rules(config.logical_axis_rules):
-      sharded_state = create_sharded_state()
+      sharded_zeros = create_sharded_zeros()
+      sharded_state = jax.tree.map(
+          lambda var, val: var.replace(value=val),
+          abstract_state,
+          sharded_zeros,
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
+      )
     model = nnx.merge(graphdef, sharded_state)
 
     # print weights sharding info under debug sharding mode
@@ -972,6 +993,27 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx):
           return target
         new_target = {}
         for k, v in target.items():
+          if k == "decoder" and isinstance(v, dict) and isinstance(meta_tree.get("decoder"), dict):
+            dec_meta = meta_tree["decoder"]
+            if "layers" in dec_meta and any(isinstance(x, str) and x.startswith("layers_") for x in v.keys()):
+              new_dec_target = {}
+              layers_dict = {}
+              layers_meta = dec_meta["layers"]
+              for dk, dv in v.items():
+                if isinstance(dk, str) and dk.startswith("layers_") and dk[7:].isdigit():
+                  idx_str = dk[7:]
+                  layer_meta = {}
+                  if isinstance(layers_meta, dict):
+                    layer_meta = layers_meta.get(idx_str, layers_meta.get(int(idx_str), {}))
+                  elif isinstance(layers_meta, (list, tuple)) and int(idx_str) < len(layers_meta):
+                    layer_meta = layers_meta[int(idx_str)]
+                  layers_dict[idx_str] = _adjust_target_for_moe_fusion(dv, layer_meta, is_nnx)
+                else:
+                  new_dec_target[dk] = _adjust_target_for_moe_fusion(dv, dec_meta.get(dk, {}), is_nnx)
+              new_dec_target["layers"] = layers_dict
+              new_target[k] = new_dec_target
+              continue
+
           if k == "wi" and "wi" not in meta_tree and "wi_0" in meta_tree and "wi_1" in meta_tree:
             if not is_nnx:
               arr = v
@@ -1110,7 +1152,7 @@ def _free_device_memory(node):
 
         return node
 
-      jax.tree_util.tree_map(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable))
+      # jax.tree_util.tree_map(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable))
 
       restored = ckptr.restore(
           epath.Path(config.load_parameters_path),
@@ -1128,6 +1170,12 @@ def _free_device_memory(node):
         )
       else:
         checkpoint = restored["params"]["params"]
+        if isinstance(checkpoint, dict) and "decoder" in checkpoint and "layers" in checkpoint["decoder"]:
+          dec = checkpoint["decoder"]
+          layers_dict = dec.pop("layers", {})
+          if isinstance(layers_dict, dict):
+            for idx_str, layer_val in layers_dict.items():
+              dec[f"layers_{idx_str}"] = layer_val
 
       if checkpoint:
         # Same QTensor caveat as `_build_value_target` / `_free_device_memory`:

tests/integration/checkpointing_test.py

@@ -232,5 +232,5 @@ def get_cmd(steps, metrics_file):
   # reports them as missing rather than leaving them at their init values.
   message = str(excinfo.value)
   assert "Checkpoint does not match the model" in message
-  assert "decoder/layers/0/" in message
+  assert "decoder/layers_0/" in message
   assert "scan_layers" in message