feat(nnx): support native Flax NNX PEFT/LoRA training loop

src/maxtext/layers/linears.py

@@ -70,6 +70,35 @@ def canonicalize_tuple(x):
     return (x,)
 
 
+def _restore_custom_types_from_state(val):
+  """Restores qwix custom quantized types from nnx.State representation."""
+  if isinstance(val, nnx.State):
+    pure_dict = {k: _restore_custom_types_from_state(v) for k, v in val.items()}
+    if "array" in pure_dict:
+      # pylint: disable=import-outside-toplevel
+      from qwix._src.providers.ptq import WithAux
+
+      return WithAux(array=pure_dict["array"], how=None)
+    elif "qvalue" in pure_dict and "scale" in pure_dict:
+      # pylint: disable=import-outside-toplevel
+      from qwix._src.core.qarray import QArray
+
+      return QArray(
+          qvalue=pure_dict["qvalue"],
+          scale=pure_dict["scale"],
+          zero_point=pure_dict.get("zero_point", None),
+          qtype=None,
+      )
+    else:
+      return pure_dict
+  elif isinstance(val, nnx.Variable):
+    return _restore_custom_types_from_state(val.get_value())
+  elif isinstance(val, dict):
+    return {k: _restore_custom_types_from_state(v) for k, v in val.items()}
+  else:
+    return val
+
+
 def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul_precision, quant):
   """Computes a dot_general operation that may be quantized."""
   dot_general = lax.dot_general
@@ -222,8 +251,16 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
     if quantizations.in_serve_mode(self.quant):
       kernel_shape = self.in_features_shape + self.out_features_shape
       kernel = jnp.zeros(kernel_shape, dtype=self.dtype)
+      restored_to_state = False
     else:
-      kernel = self.kernel[...]
+      kernel_val = self.kernel._raw_value if hasattr(self.kernel, "_raw_value") else self.kernel
+      restored_to_state = False
+      if isinstance(kernel_val, nnx.State):
+        kernel = _restore_custom_types_from_state(kernel_val)
+        self.kernel.value = kernel
+        restored_to_state = True
+      else:
+        kernel = self.kernel[...]
       # Move logit_dense kernel to device if parameter offloading is enabled
       if self.parameter_memory_host_offload:
         max_logging.log("linear.py: Moving parameter logits_dense kernel to device")
@@ -246,6 +283,9 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
         out_sharding,
     )
 
+    if restored_to_state:
+      self.kernel.value = kernel_val
+
     if self.bias is not None:
       bias = jnp.asarray(self.bias[...], self.dtype)
       output += bias

src/maxtext/trainers/pre_train/train.py

@@ -78,9 +78,7 @@
 def get_first_step(model, state):
   if isinstance(model, nn.Module):
     return int(state.step)
-  if hasattr(
-      state, "inner_state"
-  ):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
+  if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
     return int(state.step.get_value())
   return int(state.optimizer.step.get_value())
 
@@ -393,9 +391,10 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
           is_train=True,
       )
     else:
+      wrt = nnx.LoRAParam if config.lora.enable_lora else nnx.Param
       owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True)
       custom_param_filter = nnx.Any(owg_type)
-      model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, nnx.Param, custom_param_filter, ...)
+      model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, wrt, custom_param_filter, ...)
       if config.parameter_memory_host_offload:
         # Params are kept on host (pinned_host) in in_shardings. Move only Param
         # variables to device before the forward/backward pass so that all dot_general
@@ -420,7 +419,7 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
       def diff_wrapper(curr_params, custom_params, rest, config, data):
         local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
         loss, aux = loss_fn(local_model, config, data, None, None, is_train=True)
-        _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...)
+        _, _, _, new_rest = nnx.split(local_model, wrt, custom_param_filter, ...)
         return loss, (aux, new_rest)
 
       grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True)
@@ -771,6 +770,17 @@ def train_loop(config, recorder, state=None):
   start_step = get_first_step(model, state)  # this is the start_step for training
   train_utils.validate_completed_steps(start_step, config.steps)
 
+  if config.pure_nnx:
+    if isinstance(state, nnx.Module):
+      _ = nnx.pop(state, nnx.Intermediate)
+    elif isinstance(state, nnx.State):
+      state = nnx.state(state, nnx.Not(nnx.Intermediate))
+
+    if isinstance(state_mesh_shardings, nnx.Module):
+      _ = nnx.pop(state_mesh_shardings, nnx.Intermediate)
+    elif isinstance(state_mesh_shardings, nnx.State):
+      state_mesh_shardings = nnx.state(state_mesh_shardings, nnx.Not(nnx.Intermediate))
+
   if isinstance(model, nn.Module):
     jit_model = model
   elif config.enable_diloco:
@@ -784,11 +794,7 @@ def train_loop(config, recorder, state=None):
     # the Zero-1 opt overlay doesn't apply through the diloco wrapper.
     params_shardings = state_mesh_shardings.params
   else:
-    params_shardings, state_mesh_shardings = (
-        sharding.maybe_update_params_sharding_with_opt(
-            config, state_mesh_shardings
-        )
-    )
+    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config,
@@ -828,9 +834,7 @@ def train_loop(config, recorder, state=None):
   if isinstance(model, nn.Module):
     setup_params = state.params
   elif config.enable_diloco:
-    setup_params = (
-        state.params
-    )  # DiLoCoTrainState.params: the outer (global) params
+    setup_params = state.params  # DiLoCoTrainState.params: the outer (global) params
   else:
     _, setup_params, _ = nnx.split(state.model, nnx.Param, ...)
   metric_logger_instance.write_setup_info_to_tensorboard(setup_params)

src/maxtext/utils/lora_utils.py

@@ -424,14 +424,20 @@ def _get_lora_module_path(mt_config: pyconfig.HyperParameters) -> str:
   model_name = mt_config.model_name.lower()
 
   # Find the first matching architecture prefix or use 'default'
-  matched_key = next((k for k in lora_configs if k != "default" and model_name.startswith(k)), "default")
+  matched_key = next(
+      (k for k in lora_configs if k != "default" and model_name.startswith(k)),
+      "default",
+  )
 
   if matched_key == "default":
     max_logging.log(f"Warning: Model '{model_name}' is unverified; falling back to default LoRA path.")
   else:
     max_logging.log(f"Auto-detected lora_module_path for model '{model_name}' (matched: '{matched_key}')")
 
-  raw_path = lora_configs.get(matched_key, "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))")
+  raw_path = lora_configs.get(
+      matched_key,
+      "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))",
+  )
 
   # This regex makes the layer index optional, matching both scanned and unscanned layer paths
   # (e.g. 'layers/0/mlp/...' vs 'layers/mlp/...').
@@ -607,12 +613,16 @@ def apply_lora_to_model(
   )
 
   if mesh is not None:
-    with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+    with nn_partitioning.axis_rules(mt_config.logical_axis_rules):
       graph_def, state = nnx.split(lora_model)
 
       # We handle explicit replication for LoRA to ensure safety and efficiency.
       state = jax.tree_util.tree_map(
-          lambda x: x.replace(sharding=jax.sharding.PartitionSpec(), out_sharding=None, sharding_names=None)
+          lambda x: x.replace(
+              sharding=jax.sharding.PartitionSpec(),
+              out_sharding=None,
+              sharding_names=None,
+          )
           if isinstance(x, nnx.LoRAParam)
           else x,
           state,
@@ -624,18 +634,24 @@ def apply_lora_to_model(
       dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, mt_config.logical_axis_rules)
 
       def _safe_reshard(var, sharding_spec):
-        if not isinstance(var, nnx.Variable) or not isinstance(sharding_spec, jax.sharding.Sharding):
+        physical_sharding = sharding_spec.get_value() if isinstance(sharding_spec, nnx.Variable) else sharding_spec
+        if not isinstance(var, nnx.Variable) or not isinstance(physical_sharding, jax.sharding.Sharding):
           return var
         val = var.get_value()
-        if not isinstance(val, jax.Array):
+        if not isinstance(val, jax.Array) or isinstance(val, jax.core.Tracer):
           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.
-        resharded_val = jax.make_array_from_callback(val.shape, sharding_spec, lambda idx: val[idx])
+        resharded_val = jax.make_array_from_callback(val.shape, physical_sharding, lambda idx: val[idx])
         return var.replace(value=resharded_val)
 
-      state = jax.tree_util.tree_map(_safe_reshard, state, dst_shardings, is_leaf=lambda x: isinstance(x, nnx.Variable))
+      state = jax.tree_util.tree_map(
+          _safe_reshard,
+          state,
+          dst_shardings,
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
+      )
 
       lora_model = nnx.merge(graph_def, state)
 
@@ -677,7 +693,7 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter
   abstract_lora_params = nnx.state(model, nnx.LoRAParam)
 
   target_for_restore = jax.tree.map(
-      lambda v: {"value": v.value},
+      lambda v: {"value": v.get_value()},
       abstract_lora_params,
       is_leaf=lambda n: isinstance(n, nnx.Variable),
   )
@@ -699,6 +715,12 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter
     max_logging.log(f"Guided restore failed: {e}. Falling back to basic restore.")
     restored_lora_params = ocp.PyTreeCheckpointer().restore(lora_restore_path)
 
+  # If restoring from a full TrainState checkpoint, navigate into the model sub-tree
+  if isinstance(restored_lora_params, dict) and "model" in restored_lora_params:
+    restored_lora_params = restored_lora_params["model"]
+  elif hasattr(restored_lora_params, "model"):
+    restored_lora_params = getattr(restored_lora_params, "model")
+
   # Post processing
   def _map_to_state(path, variable):
     if not isinstance(variable, nnx.Variable):
@@ -722,7 +744,7 @@ def _map_to_state(path, variable):
     else:
       matched_val = curr
 
-    variable.value = matched_val
+    variable[...] = matched_val
 
   jax.tree_util.tree_map_with_path(
       _map_to_state,
@@ -839,7 +861,11 @@ def get_lora_abstract_state_nnx(base_abstract_params, lora_config):
   def get_lora_param_shape(base_array_shape, lora_module):
     if len(base_array_shape) > 4:
       raise ValueError(f"Unsupported base array shape {base_array_shape} (>4D)")
-    if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"):
+    if lora_module in (
+        "self_attention.query",
+        "self_attention.key",
+        "self_attention.value",
+    ):
       lora_a_shape = base_array_shape[:-2] + (lora_rank,)
       lora_b_shape = (lora_rank,) + base_array_shape[1:]
     elif lora_module == "self_attention.out":
@@ -858,7 +884,11 @@ def get_lora_param_sharding(base_param_sharding, lora_module):
     base_pspec = base_param_sharding.spec
     if len(base_pspec) > 4:
       raise ValueError("PartitionSpec size > 4 not supported")
-    if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"):
+    if lora_module in (
+        "self_attention.query",
+        "self_attention.key",
+        "self_attention.value",
+    ):
       lora_a_pspec = jax.sharding.PartitionSpec(*(base_pspec[:-2] + ((),)))
       lora_b_pspec = jax.sharding.PartitionSpec(*(((),) + base_pspec[1:]))
     elif lora_module == "self_attention.out":

src/maxtext/utils/sharding.py

@@ -176,9 +176,7 @@ def remove_size_one_mesh_axis(spec, mesh):
   return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced)
 
 
-def get_nnx_var_named_sharding_with_scan_axis(
-    v: nnx.Variable, mesh
-) -> nnx.Variable:
+def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable:
   """Compute NamedSharding for an NNX variable, correctly handling the scan axis."""
   val = v.get_value()
   if not hasattr(val, "shape"):
@@ -190,23 +188,15 @@ def get_nnx_var_named_sharding_with_scan_axis(
       return v.replace(jax.tree.map(lambda _: replicated, val))
     return v
   metadata = v.get_metadata()
-  out_sharding = (
-      metadata.get("out_sharding")
-      or metadata.get("sharding_names")
-      or metadata.get("sharding")
-  )
+  out_sharding = metadata.get("out_sharding") or metadata.get("sharding_names") or metadata.get("sharding")
   if not out_sharding:
     pspec = P()
   else:
     # Insert the scan axis for parameters created by _create_scanned_layers.
     if nnx.PARTITION_NAME in metadata:
       partition_name = metadata[nnx.PARTITION_NAME]
       scan_axis = metadata.get("param_scan_axis", 0)
-      out_sharding = (
-          [out_sharding]
-          if isinstance(out_sharding, str)
-          else list(out_sharding)
-      )
+      out_sharding = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding)
       if partition_name not in out_sharding:
         out_sharding.insert(scan_axis, partition_name)
       out_sharding = tuple(out_sharding)
@@ -215,9 +205,7 @@ def get_nnx_var_named_sharding_with_scan_axis(
     local_rules = metadata.get("sharding_rules", ())
     if context_rules or local_rules:
       local_rules_list = list(local_rules) if local_rules is not None else []
-      context_rules_list = (
-          list(context_rules) if context_rules is not None else []
-      )
+      context_rules_list = list(context_rules) if context_rules is not None else []
       rules = local_rules_list + context_rules_list
       pspec = logical_to_mesh_axes(out_sharding, mesh, rules=rules)
     else:
@@ -609,27 +597,33 @@ def maybe_update_params_sharding_with_opt_nnx(
   # In TrainStateNNX, parameters are under 'model'
   model_shardings = state_mesh_shardings.model
 
+  wrt = (
+      nnx.LoRAParam
+      if (hasattr(config, "lora") and config.lora and getattr(config.lora, "enable_lora", False))
+      else nnx.Param
+  )
+
   def _extract_param_only(state):
-    """Recursively extract nnx.Param variables from an nnx.State into a nested plain dict.
+    """Recursively extract trainable parameters from an nnx.State into a nested plain dict.
 
     Constructs nnx.State({'key': nested_dict, ...}) which produces the same pytree
-    structure as nnx.split(model, nnx.Param, ...)[1], enabling jax.tree.map
-    to work correctly between ga_params (Param-only) and params_shardings.
+    structure as nnx.split(model, wrt, ...)[1], enabling jax.tree.map
+    to work correctly between ga_params (wrt-only) and params_shardings.
     """
     result = {}
     for k, v in state.items():
-      if isinstance(v, nnx.Param):
+      if isinstance(v, wrt):
         result[k] = v
       elif isinstance(v, nnx.Variable):
-        pass  # skip non-Param variables (RngKey, RngCount, OptVariable, etc.)
+        pass  # skip non-trainable variables (RngKey, RngCount, OptVariable, etc.)
       elif hasattr(v, "items"):
         sub = _extract_param_only(v)
         if sub:
           result[k] = sub
     return result
 
   # prev_params_shardings must match the pytree structure of ga_params from
-  # nnx.split(model, nnx.Param, ...) — Param variables only, no rngs.
+  # nnx.split(model, wrt, ...) — trainable variables only, no rngs.
   prev_params_shardings = nnx.State(_extract_param_only(model_shardings))
 
   if not config.shard_optimizer_over_data:
@@ -701,9 +695,7 @@ def _update_model_var(path, var):
   return prev_params_shardings, updated_state
 
 
-def build_zero1_input_state_mesh_shardings(
-    config, state_mesh_shardings, params_shardings
-):
+def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_shardings):
   """Build the train-step input shardings under shard_optimizer_over_data (Zero-1).
 
   Model params on input use the original pre-Zero-1 sharding (params_shardings),