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

src/maxtext/common/checkpointing.py

@@ -196,7 +196,10 @@ def _expected_and_restored_params(abstract_nnx_state, restored_linen):
   Splits the abstract by Variable type (nnx.Param) so only real weights are compared --
   rngs/dropout/batch stats live in `nnx_aux` and are restored separately.
   """
-  want = nnx.split_state(abstract_nnx_state, nnx.Param, ...)[0].to_pure_dict().get("model", {})
+  if hasattr(nnx, "LoRAParam") and nnx.state(abstract_nnx_state, nnx.LoRAParam):
+    want = nnx.split_state(abstract_nnx_state, nnx.LoRAParam, ...)[0].to_pure_dict().get("model", {})
+  else:
+    want = nnx.split_state(abstract_nnx_state, nnx.Param, ...)[0].to_pure_dict().get("model", {})
   have = restored_linen.get("params", {}).get("params", {})
   return want, have
 
@@ -211,6 +214,7 @@ def _raise_on_weight_mismatch(want, have):
   without naming the weight.
   """
   problems = _weight_mismatches(want, have)
+
   if not problems:
     return
   lines = "\n".join(f"  - '{p}': {why}" for p, why in problems)
@@ -715,6 +719,17 @@ def _restore_grain_iterator(
   return (restored_state, None)
 
 
+def dict_has_key_substring(d, substring):
+  if not isinstance(d, dict):
+    return False
+  for k, v in d.items():
+    if isinstance(k, str) and substring in k:
+      return True
+    if isinstance(v, dict) and dict_has_key_substring(v, substring):
+      return True
+  return False
+
+
 def load_state_if_possible(
     checkpoint_manager: CheckpointManager | None,
     data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None,
@@ -888,7 +903,20 @@ def map_to_pspec(data):
     return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config)
   elif load_parameters_from_path != "":
     if isinstance(abstract_unboxed_pre_state, nnx.State):
-      _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...)
+      checkpoint_has_lora = False
+      if load_parameters_from_path:
+        try:
+          ckptr = ocp.Checkpointer(ocp.PyTreeCheckpointHandler())
+          metadata = ckptr.metadata(epath.Path(load_parameters_from_path))
+          checkpoint_has_lora = dict_has_key_substring(metadata.item_metadata.tree, "lora")
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          max_logging.log(f"Failed to read metadata from {load_parameters_from_path}: {e}")
+          checkpoint_has_lora = True
+
+      if hasattr(nnx, "LoRAParam") and not checkpoint_has_lora:
+        _, _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.LoRAParam, nnx.Param, ...)
+      else:
+        _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...)
     else:
       params = abstract_unboxed_pre_state.params
 
@@ -1086,9 +1114,14 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
       step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
       state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
     else:
-      # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout
-      # stream continues across resumes instead of resetting to a base key.
-      state = train_state_nnx.to_checkpoint_dict(state)
+      if getattr(getattr(config, "lora", None), "enable_lora", False):
+        pure_dict = state.to_pure_dict()
+        pure_dict["model"] = nnx.state(state.model, nnx.LoRAParam).to_pure_dict()
+        state = train_state_nnx.to_linen_checkpoint_dict(pure_dict)
+      else:
+        # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout
+        # stream continues across resumes instead of resetting to a base key.
+        state = train_state_nnx.to_checkpoint_dict(state)
 
   try:
     checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save)

src/maxtext/layers/linears.py

@@ -37,6 +37,7 @@
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
+from maxtext.utils import lora_utils
 from maxtext.utils.sharding import maybe_shard_with_logical
 
 
@@ -222,8 +223,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 = lora_utils.restore_qlora_base_weights(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 +255,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

@@ -391,9 +391,11 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
           is_train=True,
       )
     else:
+      enable_lora = getattr(getattr(config, "lora", None), "enable_lora", False)
+      wrt = nnx.LoRAParam if 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
@@ -418,7 +420,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)

src/maxtext/utils/lora_utils.py

@@ -26,8 +26,11 @@
 from flax.training import train_state
 import jax
 import jax.numpy as jnp
+from jax.core import Tracer
 from orbax import checkpoint as ocp
 import qwix
+from qwix._src.core.qarray import QArray
+from qwix._src.providers.ptq import WithAux
 
 from maxtext.common import checkpointing
 from maxtext.configs import pyconfig
@@ -584,7 +587,7 @@ def apply_lora_to_model(
     max_logging.log("MaxText LoRA adapters loaded, skipping Qwix LoRA application")
     return model
 
-  if not mt_config.lora.enable_lora:
+  if not getattr(getattr(mt_config, "lora", None), "enable_lora", False):
     return model
 
   # Dynamically detect and set LoRA rank before model creation if restoring
@@ -607,37 +610,38 @@ def apply_lora_to_model(
   )
 
   if mesh is not None:
-    with jax.set_mesh(mesh), 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)
-          if isinstance(x, nnx.LoRAParam)
-          else x,
-          state,
-          is_leaf=lambda x: isinstance(x, nnx.Variable),
-      )
+    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)
+        if isinstance(x, nnx.LoRAParam)
+        else x,
+        state,
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
 
-      # Use logical_to_mesh_sharding to correctly map logical axes like 'embed'
-      # to physical mesh axes.
-      dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, mt_config.logical_axis_rules)
+    # Use logical_to_mesh_sharding to correctly map logical axes like 'embed'
+    # to physical mesh axes.
+    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):
-          return var
-        val = var.get_value()
-        if not isinstance(val, jax.Array):
-          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])
-        return var.replace(value=resharded_val)
+    def _safe_reshard(var, sharding_spec):
+      if isinstance(sharding_spec, nnx.Variable):
+        sharding_spec = sharding_spec.get_value()
+      if not isinstance(var, nnx.Variable) or not isinstance(sharding_spec, jax.sharding.Sharding):
+        return var
+      val = var.get_value()
+      if not isinstance(val, jax.Array) or isinstance(val, 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])
+      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)
+    lora_model = nnx.merge(graph_def, state)
 
   _verify_lora_parameters(lora_model, mt_config)
 
@@ -666,7 +670,7 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter
 
   if not is_lora_enabled(model):
     lora_module_path = _get_lora_module_path(mt_config)
-    if not mt_config.lora.enable_lora:
+    if not getattr(getattr(mt_config, "lora", None), "enable_lora", False):
       raise ValueError(
           "lora_restore_path is set but LoRA is not enabled on the model. "
           f"Set lora.enable_lora=True and verify lora_module_path ('{lora_module_path}') matches model modules."
@@ -699,37 +703,54 @@ 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")
+  restored_count = 0
+
   # Post processing
   def _map_to_state(path, variable):
+    nonlocal restored_count
     if not isinstance(variable, nnx.Variable):
       return
 
     str_path = [str(k.key if hasattr(k, "key") else (k.name if hasattr(k, "name") else k)) for k in path]
 
     curr = restored_lora_params
     for p in str_path:
-      if isinstance(curr, dict) and p in curr:
-        curr = curr[p]
+      if isinstance(curr, Mapping):
+        if p in curr:
+          curr = curr[p]
+        elif p.isdigit() and int(p) in curr:
+          curr = curr[int(p)]
+        else:
+          return
       elif hasattr(curr, p):
         curr = getattr(curr, p)
       else:
         return
 
-    if isinstance(curr, dict) and "value" in curr:
+    if isinstance(curr, Mapping) and "value" in curr:
       matched_val = curr["value"]
     elif hasattr(curr, "value"):
       matched_val = getattr(curr, "value")
     else:
       matched_val = curr
 
     variable.value = matched_val
+    restored_count += 1
 
   jax.tree_util.tree_map_with_path(
       _map_to_state,
       abstract_lora_params,
       is_leaf=lambda n: isinstance(n, nnx.Variable),
   )
 
+  if restored_count == 0:
+    raise ValueError(f"No LoRA/adapter parameters were successfully restored from checkpoint at '{lora_restore_path}'.")
+
   nnx.update(model, abstract_lora_params)
   max_logging.log(f"LoRA restore complete from '{lora_restore_path}'.")
   return model
@@ -918,3 +939,26 @@ def add_lora(out_node, base_node, path):
       opt_state={},
   )
   return unboxed_abstract_lora_state, lora_state_mesh_annotations
+
+
+def restore_qlora_base_weights(val):
+  """Restores qwix custom quantized types from nnx.State representation."""
+  if isinstance(val, nnx.State):
+    pure_dict = {k: restore_qlora_base_weights(v) for k, v in val.items()}
+    if "array" in pure_dict:
+      return WithAux(array=pure_dict["array"], how=None)
+    elif "qvalue" in pure_dict and "scale" in pure_dict:
+      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_qlora_base_weights(val.get_value())
+  elif isinstance(val, dict):
+    return {k: restore_qlora_base_weights(v) for k, v in val.items()}
+  else:
+    return val

src/maxtext/utils/maxtext_utils.py

@@ -1688,6 +1688,20 @@ def setup_initial_state(
         # didn't carry is still an unmaterialized placeholder. Fill those from the fresh init: a
         # present leaf comes from the checkpoint, an absent one keeps its init value.
         overlay = restored if is_emergency else restored["items"]
+
+        if getattr(getattr(config, "lora", None), "enable_lora", False):
+          path_to_load = config.load_parameters_path or config.load_full_state_path
+          if path_to_load:
+            _, base_params, _ = nnx.split_state(state.model, nnx.LoRAParam, nnx.Param, ...)
+            loaded_base_params = checkpointing.load_params_from_path(
+                path_to_load,
+                base_params,
+                config.checkpoint_storage_concurrent_gb,
+                use_ocdbt=config.checkpoint_storage_use_ocdbt,
+                use_zarr3=config.checkpoint_storage_use_zarr3,
+            )
+            nnx.update(state.model, loaded_base_params)
+
         merged = jax.tree.map(
             lambda ckpt, init: init if isinstance(ckpt, jax.ShapeDtypeStruct) else ckpt,
             overlay.to_pure_dict(),
@@ -1733,6 +1747,27 @@ def _merge_params(p_raw, p_init):
             state = state.replace(params=merged_params)
           else:
             state = state.replace(params=raw_params)
+  if config.pure_nnx and getattr(getattr(config, "lora", None), "enable_lora", False):
+    # pylint: disable=import-outside-toplevel
+    from maxtext.utils import lora_utils
+
+    def _convert_var(v):
+      if isinstance(v, nnx.Variable):
+        raw_val = v.get_value() if hasattr(v, "get_value") else v.value
+        if isinstance(raw_val, nnx.State):
+          restored_val = lora_utils.restore_qlora_base_weights(raw_val)
+          if hasattr(v, "_raw_value"):
+            # pylint: disable=protected-access
+            v._raw_value = restored_val
+          v.set_value(restored_val)
+      return v
+
+    if hasattr(state, "model"):
+      model_to_convert = state.model
+      if isinstance(model_to_convert, nnx.Module):
+        model_to_convert = nnx.state(model_to_convert)
+      jax.tree_util.tree_map(_convert_var, model_to_convert, is_leaf=lambda x: isinstance(x, nnx.Variable))
+
   if not config.pure_nnx:
     state = max_utils.unbox_logicallypartioned(state)
   return state, state_mesh_annotations, state_mesh_shardings, data_iterator, was_restored