feat(nnx): support native Flax NNX PEFT/LoRA training loop with pre-trained base weight restoration and optimized upfront conversion

Author: RexBearIU

src/maxtext/common/checkpointing.py

@@ -211,6 +211,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)
@@ -1086,9 +1087,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(config, "lora", None) and getattr(config.lora, "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

@@ -28,6 +28,8 @@
 import jax.numpy as jnp
 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
@@ -607,37 +609,36 @@ 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 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)
 
-      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)
 
@@ -699,37 +700,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 +936,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(config, "lora", None) and getattr(config.lora, "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(config, "lora", None) and getattr(config.lora, "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

src/maxtext/utils/sharding.py

@@ -598,27 +598,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:

src/maxtext/utils/train_utils.py

@@ -36,6 +36,7 @@
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
 from maxtext.utils import sharding
+from maxtext.utils import lora_utils
 from maxtext.utils.rampup_batch import create_rampup_manager
 
 
@@ -241,7 +242,11 @@ def setup_train_loop(config, recorder, devices=None):
       # For NNX, the train state is wrapped in the TrainStateNNX module.
       def create_train_state_fn():
         model = _create_model_partial()
-        optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
+        wrt = nnx.Param
+        if config.lora.enable_lora:
+          model = lora_utils.apply_lora_to_model(model, mesh, config)
+          wrt = nnx.LoRAParam
+        optimizer = nnx.Optimizer(model, tx, wrt=wrt)
         return train_state_nnx.TrainStateNNX(model, optimizer)
 
       init_state_fn = create_train_state_fn
@@ -325,8 +330,6 @@ def create_train_state_fn():
         inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings)
         state_mesh_shardings = diloco.DiLoCoTrainState(
             inner_state_shardings,
-            # Match the outer params' pure-dict structure (build_diloco_state stores
-            # outer_params via to_pure_dict), so the sharding tree matches the state tree.
             state_mesh_shardings_params.to_pure_dict()  # pyrefly: ignore[missing-attribute]
             if config.pure_nnx
             else state_mesh_shardings_params,
@@ -363,6 +366,12 @@ def create_train_state_fn():
     else:
       train_state = nnx.merge(state_graphdef, state)  # pyrefly: ignore[unbound-name]
       model = train_state.model
+
+      # Restore external pre-trained LoRA adapter weights if starting a new run
+      if config.lora.enable_lora and config.lora.lora_restore_path:
+        checkpoint_step = checkpoint_manager.latest_step() if checkpoint_manager is not None else None
+        if checkpoint_step is None:
+          lora_utils.restore_lora_from_path(model, config)
   else:
     train_state = state
 
@@ -387,6 +396,13 @@ def validate_train_config(config):
   if getattr(config, "use_dpo", False):
     raise ValueError("Legacy DPO implementation in train.py is removed. Please use post-training train_dpo.py instead.")
 
+  if getattr(config, "lora", None) and getattr(config.lora, "enable_lora", False):
+    if not getattr(config, "pure_nnx", False):
+      raise ValueError(
+          "LoRA/PEFT training in the native trainer (train.py) is only"
+          " supported when running a fully NNX model (pure_nnx=True)."
+      )
+
   assert config.run_name, "Erroring out, need a real run_name"
   if config.dataset_path and not config.dataset_path.startswith("gs://"):
     max_logging.log("WARNING: 'dataset_path' might be pointing your local file system")