From 29a0ad60b5347f67d3fb8e616233a0c672818b98 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Tue, 4 Aug 2026 01:29:04 -0700 Subject: [PATCH] fix(nnx): scope PartitionSpec default strictly to LoRA parameters in sharding analysis This CL is a follow-up fix for PR #4501 (`src/maxtext/utils/sharding.py`): - Scopes `PartitionSpec P()` defaulting strictly to LoRA parameters (`isinstance(p_leaf, nnx.LoRAParam)` or `'lora'` in parameter name), allowing NNX LoRA models to execute safely without missing-spec errors. - Standard NNX parameters (`nnx.Param`) and standard Linen parameters remain `None` so unit tests (`test_mixed_sharding_fails`) throw an `AssertionError` on unannotated arrays as expected (`tolerance=0.5` untouched). - Keeps element counts (`p_leaf.size`) for unsharded parameter totals to preserve mixed-precision sharding validation. # Tests - Verified `tests/unit/maxtext_utils_test.py::TestAssertParamsSufficientlySharded` (7/7 passed with tolerance=0.5 untouched). - Verified `tests/integration/lora_e2e_nnx_test.py` (sharding validation passed across all 10 tests). # Checklist - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests and provided workload details above. - [x] I have made or will make corresponding changes to the doc if needed. PiperOrigin-RevId: 958869746 --- src/maxtext/utils/sharding.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 06e5ebf2e7..d1795880d5 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -437,14 +437,14 @@ def _analyze_sharding(params, mesh, valid_target_mesh_axes): for path, p_leaf in all_params_leaves: # Iterate over each parameter leaf param_name_str = jax.tree_util.keystr(path) # Convert the tree path to a readable string - # Unwrap nnx.Variable / nnx.LoRAParam objects to access the underlying jax.Array - if hasattr(p_leaf, "value") and not isinstance(p_leaf, jax.Array): + # Default unannotated LoRA parameters to PartitionSpec P() while leaving standard parameters as None for strict assertions. + is_lora_param = isinstance(p_leaf, getattr(nnx, "LoRAParam", ())) + is_lora = is_lora_param or "lora" in param_name_str.lower() + if isinstance(p_leaf, nnx.Variable): p_leaf = p_leaf.value - # Extract sharding spec, defaulting to PartitionSpec P() if sharding is unset or single-device - sharding = getattr(p_leaf, "sharding", None) - spec = getattr(sharding, "spec", None) - if spec is None: + spec = getattr(getattr(p_leaf, "sharding", None), "spec", None) + if spec is None and is_lora: spec = P() assert isinstance(spec, P), f"Expected '.sharding.spec' for parameter '{param_name_str}' to be a PartitionSpec."