fix(nnx): update lora regex word boundary and multi-host array resharding

Author: RexBearIU

src/maxtext/configs/post_train/lora_module_path.yml

@@ -21,7 +21,7 @@ mistral: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
 deepseek2: "decoder/(dense_layers|moe_stack)/self_attention/(query|out|wkv_a|wkv_b)|decoder/(dense_layers|moe_stack)/(mlp|shared_experts)/(wi_0|wi_1|wo)"
 gemma2: "decoder/(scanned_blocks|layers_remainder|layers)/(self_attention_local|self_attention_global)/(query|key|value|out)|decoder/(scanned_blocks|layers_remainder|layers)/(mlp_local|mlp_global)/(wi_0|wi_1|wo)"
 gemma3: "decoder/(scanned_blocks|layers_remainder|layers)/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo|gate|up|down))"
-gemma4: "decoder/((scanned_blocks|layers_remainder)/)?layers.*/.*(self_attention/(query|key|value|out)|mlp/.*(wi_0|wi_1|wo|shared_experts/(wi_0|wi_1|wo)))"
+gemma4: 'decoder/((scanned_blocks|layers_remainder)/)?layers.*/.*(self_attention/(query|key|value|out)\b|mlp/.*(wi_0|wi_1|wo|shared_experts/(wi_0|wi_1|wo)))'
 olmo3: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
 gpt3: "decoder/layers/(self_attention/(qkv_proj|out)|mlp/(wi|wo))"
 

src/maxtext/utils/lora_utils.py

@@ -634,9 +634,8 @@ def _safe_reshard(var, sharding_spec):
         return var
       if not hasattr(val, "shape"):
         return var
-      # make_array_from_callback natively constructs a globally sharded array
-      # from the local host arrays, bypassing backend-specific device_put issues
-      # on both Pathways and McJAX.
+      if isinstance(val, jax.Array):
+        return var
       resharded_val = jax.make_array_from_callback(val.shape, sharding_spec, lambda idx: val[idx])
       return var.replace(value=resharded_val)
 

src/maxtext/utils/maxtext_utils.py

@@ -1712,7 +1712,9 @@ def setup_initial_state(
       # checkpoint didn't carry (or a params-only load didn't touch) keeps its real init
       # value, so restore doesn't depend on knowing exactly what was saved.
       state = jax.jit(
-          lambda: nnx.state(init_state_partial()),  # Get state only, mapping to out_sharding structure
+          lambda: nnx.state(
+              init_state_partial(), nnx.Not(nnx.Intermediate)
+          ),  # Get state only, mapping to out_sharding structure
           in_shardings=None,
           out_shardings=state_mesh_shardings,
       )()
@@ -1911,7 +1913,7 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru
     # which has no .shape and would raise AttributeError. We handle sharding
     # ourselves via nnx_construct_named_sharding, so auto-assignment is not needed here.
     abs_model = nnx.eval_shape(nnx_init_trainstate_fn)
-    _, abs_var_state = nnx.split(abs_model)
+    abs_var_state = nnx.state(abs_model, nnx.Not(nnx.Intermediate))
     named_sharding_state = sharding.nnx_construct_named_sharding(abs_var_state, mesh)
     abstract_state = jax.tree.map(
         lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s),