From 7b8cebfe80d9cc6905d39fcc8e14b1e4c50ba22d Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Fri, 24 Jul 2026 11:21:44 +0000 Subject: [PATCH 1/7] Fix test_scan_layers_mismatch_tpu layer path assertion Update layer path assertion in test_scan_layers_mismatch_tpu from decoder/layers/0/ to decoder/layers_0/ to match the unscanned layer naming structure introduced in PR #4504. --- tests/integration/checkpointing_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/checkpointing_test.py b/tests/integration/checkpointing_test.py index 5e263fa35f..134aadd5a2 100644 --- a/tests/integration/checkpointing_test.py +++ b/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 From e9b2d34cae5f5185759a32ca99b0b0b4302a707d Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 02:33:23 +0000 Subject: [PATCH 2/7] fix(nnx): initialize sharded zero variables directly on mesh preserving NNX Variable PyTree structure --- src/maxtext/utils/maxtext_utils_nnx.py | 17 +++++++++++++---- src/maxtext/utils/model_creation_utils.py | 19 +++++++++++++------ 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 1b67c2592c..79e9f85950 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -170,13 +170,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, jax.ShapeDtypeStruct)), + ) # 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) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 7385408379..b5121267dc 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -626,16 +626,23 @@ def create_nnx_sharded_model_hybrid(config, mesh=None, devices=None, model_mode= out_shardings = nn.logical_to_mesh_sharding(specs, mesh) @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, jax.ShapeDtypeStruct)), + ) 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 From e16c4d5a41f6e1cba231ca4071afc28f10966f3a Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 02:39:40 +0000 Subject: [PATCH 3/7] fix(nnx): fix is_leaf match in create_sharded_zeros to match out_shardings raw array structure --- .dockerignore | 1 + src/maxtext/utils/maxtext_utils_nnx.py | 2 +- src/maxtext/utils/model_creation_utils.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.dockerignore b/.dockerignore index e567ea2ff6..1ef229aaae 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,2 +1,3 @@ .git maxtext_venv +.venv diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 79e9f85950..0c786c5f1c 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -174,7 +174,7 @@ 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, jax.ShapeDtypeStruct)), + is_leaf=lambda x: isinstance(x, nnx.Variable), ) # Create the model with sharded parameters. diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index b5121267dc..d3fb74ae54 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -630,7 +630,7 @@ 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, jax.ShapeDtypeStruct)), + is_leaf=lambda x: isinstance(x, nnx.Variable), ) with mesh: From 583061ec02d78e32b3d2640f655f94572241180a Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 03:04:26 +0000 Subject: [PATCH 4/7] Fix PyTree structure mismatch in create_nnx_sharded_model by stripping nnx.Variable wrappers from out_shardings --- src/dependencies/scripts/docker_upload_runner.sh | 9 ++------- src/maxtext/utils/maxtext_utils_nnx.py | 15 ++++++++++----- src/maxtext/utils/model_creation_utils.py | 5 +++++ 3 files changed, 17 insertions(+), 12 deletions(-) diff --git a/src/dependencies/scripts/docker_upload_runner.sh b/src/dependencies/scripts/docker_upload_runner.sh index f7b98a4581..d31162458c 100644 --- a/src/dependencies/scripts/docker_upload_runner.sh +++ b/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} . diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 0c786c5f1c..6fc1e14686 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/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)), ) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index d3fb74ae54..e9edcb1686 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -624,6 +624,11 @@ 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_zeros(): From 3bb8af2709b88ece4f91ed1ee7b2d074719abcd5 Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 03:11:21 +0000 Subject: [PATCH 5/7] Fix Linen checkpoint structure mapping (layers_N <-> layers[N]) when loading into NNX Gemma4 model --- src/maxtext/utils/model_creation_utils.py | 30 +++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index e9edcb1686..644566cab0 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/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 @@ -984,6 +993,21 @@ 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 = {} + for dk, dv in v.items(): + if isinstance(dk, str) and dk.startswith("layers_") and dk[7:].isdigit(): + idx_str = dk[7:] + layers_dict[idx_str] = _adjust_target_for_moe_fusion(dv, dec_meta.get("layers", {}).get(idx_str, {}), 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 @@ -1140,6 +1164,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`: From 09cc220831ce195c0311f1946503e6c5b6f26809 Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 03:20:28 +0000 Subject: [PATCH 6/7] Fix layer_meta dictionary/list lookup in _adjust_target_for_moe_fusion --- src/maxtext/utils/model_creation_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 644566cab0..662d6f825f 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -998,10 +998,16 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx): 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:] - layers_dict[idx_str] = _adjust_target_for_moe_fusion(dv, dec_meta.get("layers", {}).get(idx_str, {}), is_nnx) + 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 From f3206b2a9d842beca4dff40157e525b0afc3b50a Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Mon, 27 Jul 2026 03:27:10 +0000 Subject: [PATCH 7/7] Remove manual _free_device_memory call before restore to prevent deleted array errors for non-restored parameters --- src/maxtext/utils/model_creation_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 662d6f825f..ed5c132974 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -1152,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),