diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index da97b440ee..80c68f58ed 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/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)
diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py
index efeb085d4a..77ca1986f8 100644
--- a/src/maxtext/layers/linears.py
+++ b/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
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 8584ddcfa3..49f4ca2dc4 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/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)
diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 1aa443f093..ca299ccaff 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/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,8 +703,16 @@ 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
 
@@ -708,14 +720,19 @@ def _map_to_state(path, variable):
 
     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")
@@ -723,6 +740,7 @@ def _map_to_state(path, variable):
       matched_val = curr
 
     variable.value = matched_val
+    restored_count += 1
 
   jax.tree_util.tree_map_with_path(
       _map_to_state,
@@ -730,6 +748,9 @@ def _map_to_state(path, variable):
       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
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 92a50dc887..bf0402a824 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/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
diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py
index 80b30b14df..d423f7c144 100644
--- a/src/maxtext/utils/sharding.py
+++ b/src/maxtext/utils/sharding.py
@@ -598,19 +598,21 @@ def maybe_update_params_sharding_with_opt_nnx(
   # In TrainStateNNX, parameters are under 'model'
   model_shardings = state_mesh_shardings.model
 
+  wrt = nnx.LoRAParam if getattr(getattr(config, "lora", None), "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:
@@ -618,7 +620,7 @@ def _extract_param_only(state):
     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:
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 0b8fb6564e..03b6c76441 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/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
 
 
@@ -239,19 +240,32 @@ def setup_train_loop(config, recorder, devices=None):
 
     if config.pure_nnx:
       # For NNX, the train state is wrapped in the TrainStateNNX module.
-      def create_train_state_fn():
+      def init_state_fn_float():
         model = _create_model_partial()
         optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
         return train_state_nnx.TrainStateNNX(model, optimizer)
 
-      init_state_fn = create_train_state_fn
+      def init_state_fn_lora():
+        model = _create_model_partial()
+        model = lora_utils.apply_lora_to_model(model, mesh, config)
+        optimizer = nnx.Optimizer(model, tx, wrt=nnx.LoRAParam)
+        return train_state_nnx.TrainStateNNX(model, optimizer)
+
+      enable_lora = getattr(getattr(config, "lora", None), "enable_lora", False)
+      init_state_fn_default = init_state_fn_lora if enable_lora else init_state_fn_float
+      init_state_fn = init_state_fn_default
     else:
       init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng)
     checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn)
+    is_resuming = False
+    is_warm_start = False
     if checkpoint_manager is not None:
       checkpoint_step = checkpoint_manager.latest_step()
       if checkpoint_step is not None:
         validate_completed_steps(checkpoint_step + 1, config.steps)
+        is_resuming = True
+    if not is_resuming and config.load_parameters_path != "":
+      is_warm_start = True
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
     data_iterator, eval_data_iterator = create_data_iterator(config, mesh)
@@ -299,9 +313,34 @@ def create_train_state_fn():
     # Create data_loader AFTER reordering wrapper is applied
     data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager)
 
+    init_state_fn_to_use = init_state_fn
+    if config.pure_nnx and is_warm_start and enable_lora:
+      init_state_fn_to_use = init_state_fn_float
+
     state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
-        data_iterator, config, mesh, checkpoint_manager, init_state_fn
+        data_iterator, config, mesh, checkpoint_manager, init_state_fn_to_use
     )
+
+    if config.pure_nnx and is_warm_start and enable_lora:
+      # Post-process state to apply LoRA
+      # We need to merge to Module first to apply LoRA.
+      state_graphdef_float = nnx.graphdef(nnx.eval_shape(init_state_fn_float))
+      train_state_float = nnx.merge(state_graphdef_float, state)
+
+      model = lora_utils.apply_lora_to_model(train_state_float.model, mesh, config)
+      optimizer = nnx.Optimizer(model, tx, wrt=nnx.LoRAParam)
+      train_state_lora = train_state_nnx.TrainStateNNX(model, optimizer)
+
+      # Split back to State
+      _, state = nnx.split(train_state_lora)
+
+      # Update state_mesh_shardings
+      _, _, state_mesh_shardings = maxtext_utils.get_abstract_state(
+          config, mesh, init_state_fn_lora, is_training=True
+      )
+
+      # Update init_state_fn for downstream use (graphdef)
+      init_state_fn = init_state_fn_lora
     if config.pure_nnx:
       with nn_partitioning.axis_rules(config.logical_axis_rules):
         # We only need the graphdef here; it's merged with state below. Avoid
@@ -325,8 +364,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 +400,14 @@ 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 getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(
+          getattr(config, "lora", None), "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 +432,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(getattr(config, "lora", None), "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")
diff --git a/tests/integration/peft_integration_test.py b/tests/integration/peft_integration_test.py
new file mode 100644
index 0000000000..56bc617a44
--- /dev/null
+++ b/tests/integration/peft_integration_test.py
@@ -0,0 +1,375 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Integration tests for Flax NNX PEFT, LoRA, and QLoRA."""
+
+import os
+import sys
+import unittest
+import tempfile
+import json
+import shutil
+import pytest
+import subprocess
+
+from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
+from tests.utils.test_helpers import get_test_config_path
+
+
+def _get_jax_backend():
+  """Determine JAX backend by running a quick subprocess to avoid locking TPU in the parent process."""
+  try:
+    result = subprocess.run(
+        [sys.executable, "-c", "import jax; print(jax.default_backend())"], capture_output=True, text=True, check=True
+    )
+    return result.stdout.strip()
+  except Exception:  # pylint: disable=broad-exception-caught
+    return "cpu"
+
+
+def _get_jax_device_count():
+  """Determine JAX device count by running a quick subprocess."""
+  try:
+    result = subprocess.run(
+        [sys.executable, "-c", "import jax; print(jax.device_count())"], capture_output=True, text=True, check=True
+    )
+    return int(result.stdout.strip())
+  except Exception:  # pylint: disable=broad-exception-caught
+    return 1
+
+
+def _get_common_args(temp_dir, model_name, steps=1, overrides=None):
+  """Build common arguments for downscaled integration runs."""
+  tokenizer_path = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "tokenizer.llama2")
+
+  backend = _get_jax_backend()
+  num_devices = _get_jax_device_count()
+
+  hardware = "cpu"
+  if backend == "tpu":
+    hardware = "tpu"
+  elif backend == "gpu":
+    hardware = "gpu"
+
+  args = [
+      None,
+      get_test_config_path(),
+      f"run_name=test_peft_{model_name}",
+      f"model_name={model_name}",
+      f"steps={steps}",
+      f"base_output_directory={temp_dir}",
+      "dataset_type=synthetic",
+      "per_device_batch_size=1",
+      "max_target_length=32",
+      f"tokenizer_path={tokenizer_path}",
+      "enable_goodput_recording=False",
+      "enable_checkpoint_cloud_logger=False",
+      "monitor_goodput=False",
+      "sharding_tolerance=1.0",  # Bypass sharding checks for downscaled test models
+      f"hardware={hardware}",
+      "skip_jax_distributed_system=True",
+      "ici_fsdp_parallelism=1",
+      f"ici_data_parallelism={num_devices}",
+  ]
+  if overrides:
+    args.extend(overrides)
+  return args
+
+
+def _run_train_subprocess(args):
+  """Run the pre-training trainer inside an isolated Python subprocess."""
+  cmd = [sys.executable, "-m", "maxtext.trainers.pre_train.train", args[1]] + args[2:]
+  env = os.environ.copy()
+  env["PYTHONPATH"] = os.path.abspath("src")
+
+  backend = _get_jax_backend()
+
+  if backend not in ("tpu", "gpu"):
+    env["JAX_PLATFORMS"] = "cpu"
+
+  subprocess.run(cmd, env=env, check=True)
+
+
+def _run_sft_subprocess(args):
+  """Run the SFT native trainer inside an isolated Python subprocess."""
+  cmd = [sys.executable, "-m", "maxtext.trainers.post_train.sft.train_sft_native", args[1]] + args[2:]
+  env = os.environ.copy()
+  env["PYTHONPATH"] = os.path.abspath("src")
+
+  backend = _get_jax_backend()
+
+  if backend not in ("tpu", "gpu"):
+    env["JAX_PLATFORMS"] = "cpu"
+
+  subprocess.run(cmd, env=env, check=True)
+
+
+@pytest.mark.integration_test
+class PEFTIntegrationTest(unittest.TestCase):
+  """Integration checks for pre-training and native SFT with LoRA & QLoRA."""
+
+  def setUp(self):
+    self.temp_dirs = []
+
+  def tearDown(self):
+    for temp_dir in self.temp_dirs:
+      if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+
+  def _create_temp_dir(self):
+    temp_dir = tempfile.mkdtemp()
+    self.temp_dirs.append(temp_dir)
+    return temp_dir
+
+  def test_native_sft_gemma4_lora(self):
+    """Test native SFT with Gemma4-e2b and LoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=False",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+        "num_kv_shared_layers=1",
+    ]
+    args = _get_common_args(temp_dir, "gemma4-e2b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_gemma4_qlora(self):
+    """Test native SFT with Gemma4-e2b and QLoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=False",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+        "num_kv_shared_layers=1",
+    ]
+    args = _get_common_args(temp_dir, "gemma4-e2b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_qwen3_lora(self):
+    """Test native SFT with Qwen3-0.6B and LoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=True",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+    ]
+    args = _get_common_args(temp_dir, "qwen3-0.6b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_qwen3_qlora(self):
+    """Test native SFT with Qwen3-0.6B and QLoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=True",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+    ]
+    args = _get_common_args(temp_dir, "qwen3-0.6b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_pretrain_lora_checkpoint_save_and_restore(self):
+    """Test pre-training with LoRA, checkpoint saving, and subsequent restoration."""
+    temp_dir = self._create_temp_dir()
+    metrics_file_saved = os.path.join(temp_dir, "saved_metrics.txt")
+    metrics_file_restored = os.path.join(temp_dir, "restored_metrics.txt")
+
+    # Common model scaling/configs
+    overrides_base = [
+        "scan_layers=True",
+        "attention=dot_product",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=128",
+        "base_mlp_dim=256",
+        "ici_fsdp_parallelism=1",
+    ]
+
+    # 1. First run: train 1 step with checkpointing enabled
+    args_save = _get_common_args(
+        temp_dir,
+        "default",
+        steps=1,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"metrics_file={metrics_file_saved}",
+        ],
+    )
+    _run_train_subprocess(args_save)
+
+    # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items)
+    checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", "0", "items")
+    self.assertTrue(os.path.exists(checkpoint_dir), f"Checkpoint directory {checkpoint_dir} was not created.")
+
+    # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2)
+    args_restore = _get_common_args(
+        temp_dir,
+        "default",
+        steps=2,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"load_parameters_path={checkpoint_dir}",
+            f"metrics_file={metrics_file_restored}",
+        ],
+    )
+    _run_train_subprocess(args_restore)
+
+    # 3. Verify losses are logged and correct
+    self.assertTrue(os.path.exists(metrics_file_saved))
+    self.assertTrue(os.path.exists(metrics_file_restored))
+
+    with open(metrics_file_saved, "r", encoding="utf8") as f:
+      saved_metrics = [json.loads(line) for line in f.readlines()]
+    with open(metrics_file_restored, "r", encoding="utf8") as f:
+      restored_metrics = [json.loads(line) for line in f.readlines()]
+
+    print("SAVED METRICS:", saved_metrics)
+    print("RESTORED METRICS:", restored_metrics)
+
+    self.assertEqual(len(saved_metrics), 1, "First run should have exactly 1 step.")
+    self.assertEqual(len(restored_metrics), 1, "Second run should have exactly 1 step.")
+
+    saved_step0_loss = saved_metrics[0]["learning/loss"]
+    restored_step1_loss = restored_metrics[0]["learning/loss"]
+
+    print(f"Saved run Step 0 Loss: {saved_step0_loss}")
+    print(f"Restored run Step 1 Loss: {restored_step1_loss}")
+
+    # Verify that the restored run resumed from step 0 (logged step 1)
+    self.assertEqual(saved_metrics[0]["step"], 0.0)
+    self.assertEqual(restored_metrics[0]["step"], 1.0)
+
+  def test_pretrain_qlora_checkpoint_save_and_restore(self):
+    """Test pre-training with QLoRA, checkpoint saving, and subsequent restoration."""
+    temp_dir = self._create_temp_dir()
+    metrics_file_saved = os.path.join(temp_dir, "saved_metrics_qlora.txt")
+    metrics_file_restored = os.path.join(temp_dir, "restored_metrics_qlora.txt")
+
+    # Common model scaling/configs
+    overrides_base = [
+        "scan_layers=True",
+        "attention=dot_product",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=128",
+        "base_mlp_dim=256",
+        "ici_fsdp_parallelism=1",
+    ]
+
+    # 1. First run: train 1 step with checkpointing enabled
+    args_save = _get_common_args(
+        temp_dir,
+        "default",
+        steps=1,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"metrics_file={metrics_file_saved}",
+        ],
+    )
+    _run_train_subprocess(args_save)
+
+    # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items)
+    checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", "0", "items")
+    self.assertTrue(os.path.exists(checkpoint_dir), f"Checkpoint directory {checkpoint_dir} was not created.")
+
+    # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2)
+    args_restore = _get_common_args(
+        temp_dir,
+        "default",
+        steps=2,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"load_parameters_path={checkpoint_dir}",
+            f"metrics_file={metrics_file_restored}",
+        ],
+    )
+    _run_train_subprocess(args_restore)
+
+    # 3. Verify losses are logged and correct
+    self.assertTrue(os.path.exists(metrics_file_saved))
+    self.assertTrue(os.path.exists(metrics_file_restored))
+
+    with open(metrics_file_saved, "r", encoding="utf8") as f:
+      saved_metrics = [json.loads(line) for line in f.readlines()]
+    with open(metrics_file_restored, "r", encoding="utf8") as f:
+      restored_metrics = [json.loads(line) for line in f.readlines()]
+
+    print("SAVED METRICS (QLoRA):", saved_metrics)
+    print("RESTORED METRICS (QLoRA):", restored_metrics)
+
+    self.assertEqual(len(saved_metrics), 1, "First run should have exactly 1 step.")
+    self.assertEqual(len(restored_metrics), 1, "Second run should have exactly 1 step.")
+
+    saved_step0_loss = saved_metrics[0]["learning/loss"]
+    restored_step1_loss = restored_metrics[0]["learning/loss"]
+
+    print(f"Saved run Step 0 Loss (QLoRA): {saved_step0_loss}")
+    print(f"Restored run Step 1 Loss (QLoRA): {restored_step1_loss}")
+
+    # Verify that the restored run resumed from step 0 (logged step 1)
+    self.assertEqual(saved_metrics[0]["step"], 0.0)
+    self.assertEqual(restored_metrics[0]["step"], 1.0)
diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py
index f446c8f9df..172f771cb4 100644
--- a/tests/integration/setup_train_loop_nnx_test.py
+++ b/tests/integration/setup_train_loop_nnx_test.py
@@ -27,13 +27,15 @@
 import unittest
 
 from flax import nnx
-import jax
+import jax.numpy as jnp
+import pytest
 from maxtext.common import train_state_nnx
 from maxtext.configs import pyconfig
+from maxtext.trainers.pre_train import train
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
+from maxtext.utils.lora_utils import is_lora_enabled
 from maxtext.utils.train_utils import setup_train_loop
 from tests.utils.test_helpers import get_test_config_path
-import pytest
 
 
 def _tiny_nnx_pyconfig(**overrides):
@@ -44,6 +46,7 @@ def _tiny_nnx_pyconfig(**overrides):
       "dataset_type": "synthetic",
       "model_name": "default",
       "pure_nnx": True,
+      "attention": "dot_product",
       "per_device_batch_size": 1.0,
       "base_emb_dim": 8,
       "base_num_query_heads": 4,
@@ -122,13 +125,124 @@ def test_pure_nnx_setup_param_only_split_matches_model(self):
 
     # Same key-set after nnx.split — this is what setup_train_loop relies on at
     # train_utils.py:281-282 to pair state_params with state_mesh_shardings_params.
-    self.assertEqual(
-        jax.tree_util.tree_structure(params),
-        jax.tree_util.tree_structure(params_shardings),
+    self.assertEqual(params.keys(), params_shardings.keys())
+    del model
+
+
+class SetupTrainLoopNNXLoraTest(unittest.TestCase):
+  """Integration checks for setup_train_loop and training step with LoRA enabled."""
+
+  def test_pure_nnx_setup_with_lora_enabled(self):
+    # Overriding configurations to enable lora on default model
+    config = _tiny_nnx_pyconfig(
+        weight_dtype="bfloat16",
+        sharding_tolerance=1.0,
+        lora={
+            "enable_lora": True,
+            "lora_rank": 4,
+            "lora_alpha": 8.0,
+        },
     )
-    self.assertGreater(len(jax.tree.leaves(params)), 0)
+    (
+        _,
+        _,
+        _,
+        model,
+        _,
+        _,
+        _,
+        _,
+        _,
+        _,
+        train_state,
+    ) = setup_train_loop(config, recorder=None)
 
-    del model
+    self.assertTrue(is_lora_enabled(model))
+    self.assertIs(train_state.optimizer.wrt, nnx.LoRAParam)
+
+  def _run_train_step_updates_only_adapters_test(self, qlora: bool = False, scan_layers: bool = False):
+    lora_dict = {
+        "enable_lora": True,
+        "lora_rank": 4,
+        "lora_alpha": 8.0,
+    }
+    if qlora:
+      lora_dict.update(
+          {
+              "lora_weight_qtype": "int8",
+              "lora_tile_size": 32,
+          }
+      )
+
+    num_layers = 8 if scan_layers else 2
+    config = _tiny_nnx_pyconfig(
+        weight_dtype="bfloat16",
+        sharding_tolerance=1.0,
+        base_emb_dim=32,
+        base_mlp_dim=128,
+        base_num_decoder_layers=num_layers,
+        lora=lora_dict,
+        steps=1,
+        scan_layers=scan_layers,
+    )
+    (
+        _,
+        _,
+        _,
+        model,
+        _,
+        _,
+        _,
+        data_loader,
+        rampup_manager,
+        _,
+        train_state,
+    ) = setup_train_loop(config, recorder=None)
+
+    # Extract initial weights
+    initial_model_state = nnx.state(model)
+
+    # Run one training step
+    batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
+
+    new_state, _ = train.train_step(
+        model=nnx.graphdef(train_state),
+        config=config,
+        state_mesh_shardings=None,
+        params_shardings=None,
+        state=nnx.state(train_state),
+        data=batch,
+    )
+
+    # Merge and compare
+    new_train_state = nnx.merge(nnx.graphdef(train_state), new_state)
+    new_model_state = nnx.state(new_train_state.model)
+
+    # Verify base weights are identical, and LoRA weights are updated
+    for path, initial_var in initial_model_state.items():
+      new_var = new_model_state[path]
+      if isinstance(initial_var, nnx.LoRAParam):
+        # Should have changed
+        self.assertFalse(jnp.allclose(initial_var.value, new_var.value))
+      elif isinstance(initial_var, nnx.Param):
+        # Should NOT have changed
+        self.assertTrue(jnp.allclose(initial_var.value, new_var.value))
+
+  def test_train_step_updates_only_lora_weights(self):
+    """Test standard LoRA updates only the adapter weights during a training step."""
+    self._run_train_step_updates_only_adapters_test(qlora=False, scan_layers=False)
+
+  def test_train_step_updates_only_qlora_weights(self):
+    """Test QLoRA updates only the adapter weights during a training step."""
+    self._run_train_step_updates_only_adapters_test(qlora=True, scan_layers=False)
+
+  def test_train_step_updates_only_lora_weights_scan_layers(self):
+    """Test standard LoRA updates only the adapter weights under sequentially scanned layers."""
+    self._run_train_step_updates_only_adapters_test(qlora=False, scan_layers=True)
+
+  def test_train_step_updates_only_qlora_weights_scan_layers(self):
+    """Test QLoRA updates only the adapter weights under sequentially scanned layers."""
+    self._run_train_step_updates_only_adapters_test(qlora=True, scan_layers=True)
 
 
 if __name__ == "__main__":
diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py
index e910995542..e83677bee5 100644
--- a/tests/post_training/unit/lora_utils_test.py
+++ b/tests/post_training/unit/lora_utils_test.py
@@ -22,9 +22,11 @@
 from etils import epath
 import jax
 import jax.numpy as jnp
+import numpy as np
 import optax
 import pytest
 from flax import nnx
+from qwix._src.core.qarray import QArray
 
 # Skip the entire test suite if dependencies are missing
 pytestmark = [pytest.mark.post_training]
@@ -489,6 +491,633 @@ def test_gemma4_lora_path_matching(self):
     for path in non_matching_paths:
       self.assertFalse(compiled.search(path), f"Incorrectly matched invalid path: {path}")
 
+  def test_checkpoint_saving_lora_only(self):
+    """Test that the native checkpointer only saves LoRA parameters when enable_lora is True."""
+    # Create a mock config with lora enabled
+    cfg = _make_config(
+        lora={
+            "enable_lora": True,
+            "lora_rank": 8,
+            "lora_alpha": 16.0,
+        }
+    )
+    cfg.pure_nnx = True
+    cfg.enable_diloco = False
+    cfg.enable_checkpointing = True
+
+    # Create a mock state
+    mock_state = mock.MagicMock()
+    mock_state.to_pure_dict.return_value = {
+        "model": {
+            "decoder": {
+                "layers": {
+                    "0": {
+                        "self_attention": {
+                            "query": {
+                                "kernel": jnp.array([1.0]),  # Base weight
+                                "lora_a": jnp.array([2.0]),  # LoRA weight
+                            }
+                        }
+                    }
+                }
+            }
+        },
+        "optimizer": {
+            "step": jnp.array(0, dtype=jnp.uint32),
+            "opt_state": {},
+        },
+    }
+
+    # When nnx.state(state.model, nnx.LoRAParam) is called, return only the lora variables
+    mock_lora_state = mock.MagicMock()
+    mock_lora_state.to_pure_dict.return_value = {
+        "decoder": {
+            "layers": {
+                "0": {
+                    "self_attention": {
+                        "query": {
+                            "lora_a": jnp.array([2.0]),
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    with (
+        mock.patch("flax.nnx.state") as mock_nnx_state,
+        mock.patch("maxtext.common.checkpointing.save_checkpoint") as mock_save,
+    ):
+      mock_nnx_state.return_value = mock_lora_state
+
+      # Call maybe_save_checkpoint
+      mock_manager = mock.MagicMock()
+      mock_manager.latest_step.return_value = -1
+      mock_manager.reached_preemption.return_value = False
+
+      checkpointing.maybe_save_checkpoint(
+          checkpoint_manager=mock_manager,
+          state=mock_state,
+          config=cfg,
+          data_iterator=None,
+          step=0,
+      )
+
+      # Ensure save_checkpoint was called
+      self.assertTrue(mock_save.called)
+
+      # Extract the saved state passed to save_checkpoint
+      saved_state = mock_save.call_args[0][2]
+
+      # Verify that saved_state has the legacy Linen format ("params" and "opt_state")
+      self.assertIn("params", saved_state)
+      self.assertIn("step", saved_state)
+
+      # Verify that base weights are GONE, but LoRA parameters are SAVED
+      saved_params = saved_state["params"]["params"]
+      # Inside saved_params, only lora_a exists, kernel is gone!
+      query_params = saved_params["decoder"]["layers"]["0"]["self_attention"]["query"]
+      self.assertIn("lora_a", query_params)
+      self.assertNotIn("kernel", query_params)
+
+  def test_checkpoint_restoration_preserves_base_weights(self):
+    """Test that checkpoint restoration preserves base model weights under LoRA."""
+    # Simulating Orbax restoration behavior: missing base weights are restored as ShapeDtypeStructs
+    overlay = {
+        "model": {
+            "decoder": {
+                "layers": {
+                    "0": {
+                        "self_attention": {
+                            "query": {
+                                "kernel": jax.ShapeDtypeStruct(
+                                    (1,), jnp.float32
+                                ),  # Missing base weight is returned as ShapeDtypeStruct
+                                "lora_a": jnp.array([9.0]),  # Saved LoRA weight
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    # Live concrete state before restoration (containing pre-trained base weight [5.0] and initialized lora_a [1.0])
+    state = {
+        "model": {
+            "decoder": {
+                "layers": {
+                    "0": {
+                        "self_attention": {
+                            "query": {
+                                "kernel": jnp.array([5.0]),  # Loaded pre-trained base weight
+                                "lora_a": jnp.array([1.0]),  # Initialized LoRA weight
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    # Perform the exact same jax.tree.map merge as in setup_initial_state
+    merged = jax.tree.map(
+        lambda ckpt, init: init if isinstance(ckpt, jax.ShapeDtypeStruct) else ckpt,
+        overlay,
+        state,
+        is_leaf=lambda x: isinstance(x, jax.ShapeDtypeStruct),
+    )
+
+    query_restored = merged["model"]["decoder"]["layers"]["0"]["self_attention"]["query"]
+
+    # 1. Base weights ("kernel") must be preserved exactly as [5.0]
+    self.assertEqual(query_restored["kernel"][0], 5.0)
+
+    # 2. LoRA weights ("lora_a") must be restored exactly as [9.0]
+    self.assertEqual(query_restored["lora_a"][0], 9.0)
+
+
+# ---------------------------------------------------------------------------
+# NNX-specific LoRA Helpers & Tests
+# ---------------------------------------------------------------------------
+
+from maxtext.utils import sharding
+from maxtext.utils.lora_utils import (
+    apply_lora_on_base_params,
+    apply_lora_on_base_params_nnx,
+    get_lora_abstract_state_nnx,
+    unapply_lora_from_base_params,
+    unapply_lora_from_base_params_nnx,
+)
+
+
+def _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4, dtype=jnp.float32):
+  """Tiny NNX-shaped abstract state for one attention block."""
+
+  def _sds(shape):
+    return jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=None)
+
+  return {
+      "decoder": {
+          "layers": {
+              "self_attention": {
+                  "query": {"kernel": _sds((emb, num_heads, head_dim))},
+                  "key": {"kernel": _sds((emb, num_heads, head_dim))},
+                  "value": {"kernel": _sds((emb, num_heads, head_dim))},
+                  "out": {"kernel": _sds((emb, num_heads, head_dim))},
+              },
+              "mlp": {"wi": {"kernel": _sds((emb, 4 * emb))}},
+          },
+          "shared_embedding": {"embedding": _sds((100, emb))},
+      },
+  }
+
+
+def _make_linen_attention_abstract(emb=8, num_heads=2, head_dim=4, dtype=jnp.float32):
+  """Linen-shaped equivalent (with the `{"params": ...}` outer wrap)."""
+  return {"params": _make_nnx_attention_abstract(emb, num_heads, head_dim, dtype)}
+
+
+def _lora_config(rank=4, alpha=8.0, target_modules=("q_proj", "v_proj")):
+  return {
+      "r": rank,
+      "lora_alpha": alpha,
+      "target_modules": list(target_modules),
+  }
+
+
+class TestGetLoraAbstractStateNnx(unittest.TestCase):
+  """`get_lora_abstract_state_nnx` shape, sharding, and error-path coverage."""
+
+  def test_lora_shapes_for_query_and_value(self):
+    abs_params = _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4)
+    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(rank=4))
+    attn = state.params["decoder"]["layers"]["self_attention"]
+
+    a = attn["query"]["lora_a.kernel"]
+    b = attn["query"]["lora_b.kernel"]
+    self.assertEqual(a.shape, (8, 4))
+    self.assertEqual(b.shape, (4, 2, 4))
+    self.assertEqual(a.dtype, jnp.float32)
+    self.assertEqual(b.dtype, jnp.float32)
+
+    a = attn["value"]["lora_a.kernel"]
+    b = attn["value"]["lora_b.kernel"]
+    self.assertEqual(a.shape, (8, 4))
+    self.assertEqual(b.shape, (4, 2, 4))
+
+  def test_non_target_modules_emit_none_leaves(self):
+    abs_params = _make_nnx_attention_abstract()
+    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(target_modules=("q_proj",)))
+    attn = state.params["decoder"]["layers"]["self_attention"]
+    self.assertIn("lora_a.kernel", attn["query"])
+    self.assertIsNone(attn["key"]["kernel"])
+    self.assertIsNone(attn["value"]["kernel"])
+    self.assertIsNone(attn["out"]["kernel"])
+    self.assertIsNone(state.params["decoder"]["layers"]["mlp"]["wi"]["kernel"])
+    self.assertIsNone(state.params["decoder"]["shared_embedding"]["embedding"])
+
+  def test_o_proj_has_distinct_shape(self):
+    abs_params = _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4)
+    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(rank=3, target_modules=("o_proj",)))
+    out = state.params["decoder"]["layers"]["self_attention"]["out"]
+    a = out["lora_a.kernel"]
+    b = out["lora_b.kernel"]
+    # For a 3D base (emb, num_heads, head_dim): lora_a.shape ends with rank,
+    # lora_b shape is (rank, head_dim).
+    self.assertEqual(a.shape, (8, 2, 3))
+    self.assertEqual(b.shape, (3, 4))
+
+  def test_unsupported_leaf_type_raises(self):
+    bad = {"decoder": {"layers": {"self_attention": {"query": {"kernel": jnp.zeros((4, 2, 2))}}}}}
+    with self.assertRaises(ValueError):
+      get_lora_abstract_state_nnx(bad, _lora_config())
+
+  def test_unexpected_leaf_name_raises(self):
+    bad = {"decoder": {"layers": {"self_attention": {"query": {"weight": jax.ShapeDtypeStruct((4, 2), jnp.float32)}}}}}
+    with self.assertRaises(ValueError):
+      get_lora_abstract_state_nnx(bad, _lora_config())
+
+
+def _concrete_base(rng=None, emb=4, num_heads=2, head_dim=3):
+  """Concrete arrays mirroring the abstract structure used above (NNX-shape)."""
+  if rng is None:
+    rng = jax.random.key(0)
+  k1, k2, k3, k4, k5, k6 = jax.random.split(rng, 6)
+  shape_attn = (emb, num_heads, head_dim)
+  return {
+      "decoder": {
+          "layers": {
+              "self_attention": {
+                  "query": {"kernel": jax.random.normal(k1, shape_attn)},
+                  "key": {"kernel": jax.random.normal(k2, shape_attn)},
+                  "value": {"kernel": jax.random.normal(k3, shape_attn)},
+                  "out": {"kernel": jax.random.normal(k4, shape_attn)},
+              },
+              "mlp": {"wi": {"kernel": jax.random.normal(k5, (emb, 4 * emb))}},
+          },
+          "shared_embedding": {"embedding": jax.random.normal(k6, (100, emb))},
+      },
+  }
+
+
+def _build_lora_params(base, lora_config_dict, rng):
+  """Build a concrete LoRA tree (random arrays) matching `base`."""
+  abs_tree = jax.tree_util.tree_map(lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=None), base)
+  lora_state, _ = get_lora_abstract_state_nnx(abs_tree, lora_config_dict)
+
+  def _to_concrete(leaf, rng_key):
+    if leaf is None:
+      return None
+    return jax.random.normal(rng_key, leaf.shape, leaf.dtype)
+
+  leaves, tree = jax.tree_util.tree_flatten(lora_state.params, is_leaf=lambda x: x is None)
+  rngs = jax.random.split(rng, max(1, len(leaves)))
+  out_leaves = [_to_concrete(l, r) for l, r in zip(leaves, rngs)]
+  return jax.tree_util.tree_unflatten(tree, out_leaves)
+
+
+class TestApplyLoraNnx(unittest.TestCase):
+  """`apply_lora_on_base_params_nnx` round-trip and Linen-vs-NNX parity."""
+
+  def test_apply_then_unapply_is_identity(self):
+    rng = jax.random.key(42)
+    base_orig = _concrete_base(rng)
+    base = jax.tree_util.tree_map(jnp.copy, base_orig)
+    lora = _build_lora_params(
+        base,
+        _lora_config(rank=2, target_modules=("q_proj", "v_proj")),
+        jax.random.key(7),
+    )
+    apply_lora_on_base_params_nnx(base, lora, lora_scale_factor=0.5)
+    # The query and value kernels are targets and must have changed.
+    self.assertFalse(
+        jnp.allclose(
+            base["decoder"]["layers"]["self_attention"]["query"]["kernel"],
+            base_orig["decoder"]["layers"]["self_attention"]["query"]["kernel"],
+        )
+    )
+    # The key and out kernels are non-targets and must be untouched.
+    np.testing.assert_array_equal(
+        np.asarray(base["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
+        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
+    )
+    np.testing.assert_array_equal(
+        np.asarray(base["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
+        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
+    )
+    unapply_lora_from_base_params_nnx(base, lora, lora_scale_factor=0.5)
+    np.testing.assert_allclose(
+        np.asarray(base["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+    np.testing.assert_allclose(
+        np.asarray(base["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
+        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+  def test_numerical_parity_with_linen_apply(self):
+    """The NNX and Linen apply paths produce identical results on the same inputs."""
+    rng = jax.random.key(123)
+    base_nnx = _concrete_base(rng)
+    base_linen = {"params": jax.tree_util.tree_map(jnp.copy, base_nnx)}
+    lora = _build_lora_params(
+        base_nnx,
+        _lora_config(rank=2, target_modules=("q_proj",)),
+        jax.random.key(5),
+    )
+    apply_lora_on_base_params_nnx(base_nnx, lora, lora_scale_factor=0.7)
+    apply_lora_on_base_params(base_linen, {"params": lora}, lora_scale_factor=0.7)
+    np.testing.assert_allclose(
+        np.asarray(base_nnx["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        np.asarray(base_linen["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        rtol=1e-6,
+    )
+
+  def test_apply_with_unexpected_lora_key_raises(self):
+    base = _concrete_base()
+    bad = {"decoder": {"layers": {"self_attention": {"query": {"unexpected": jnp.zeros((4, 2))}}}}}
+    with self.assertRaises(ValueError):
+      apply_lora_on_base_params_nnx(base, bad)
+
+
+class TestLinenLoraRegression(unittest.TestCase):
+  """Smoke tests for the Linen apply / unapply helpers (no other unit test exercises them)."""
+
+  def _linen_pair(self, rng=None):
+    """Build a Linen-shape (with `{"params": ...}` outer wrapper) base + lora pair."""
+    if rng is None:
+      rng = jax.random.key(99)
+    base_inner = _concrete_base(rng)
+    base = {"params": jax.tree_util.tree_map(jnp.copy, base_inner)}
+    lora_inner = _build_lora_params(
+        base_inner,
+        _lora_config(rank=2, target_modules=("q_proj", "v_proj")),
+        jax.random.key(7),
+    )
+    lora = {"params": lora_inner}
+    return base, lora
+
+  def test_linen_apply_then_unapply_is_identity(self):
+    base, lora = self._linen_pair()
+    base_orig = jax.tree_util.tree_map(jnp.copy, base)
+    apply_lora_on_base_params(base, lora, lora_scale_factor=0.5)
+    unapply_lora_from_base_params(base, lora, lora_scale_factor=0.5)
+    np.testing.assert_allclose(
+        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+    np.testing.assert_allclose(
+        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
+        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+  def test_linen_apply_only_modifies_target_modules(self):
+    base, lora = self._linen_pair()
+    base_orig = jax.tree_util.tree_map(jnp.copy, base)
+    apply_lora_on_base_params(base, lora, lora_scale_factor=1.0)
+    # query and value are targets and must change.
+    self.assertFalse(
+        jnp.allclose(
+            base["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"],
+            base_orig["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"],
+        )
+    )
+    # key and out are non-target and must be untouched.
+    np.testing.assert_array_equal(
+        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
+        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
+    )
+    np.testing.assert_array_equal(
+        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
+        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
+    )
+
+
+class TestShardingExtractionNnx(unittest.TestCase):
+  """Test sharding parameter extraction under different LoRA configurations."""
+
+  def test_sharding_extracts_only_lora_params(self):
+    class ToyModel(nnx.Module):
+
+      def __init__(self):
+        self.p = nnx.Param(jnp.ones((2, 2)))
+        self.lora_p = nnx.LoRAParam(jnp.zeros((2, 2)))
+
+    class DummyConfig:
+
+      class DummyLora:
+        enable_lora = True
+
+      lora = DummyLora()
+      shard_optimizer_over_data = False
+      pure_nnx = True
+
+    model = ToyModel()
+    _, state = nnx.split(model)
+
+    class DummyStateMeshShardings:
+      model = state
+
+    prev_params, _ = sharding.maybe_update_params_sharding_with_opt(DummyConfig(), DummyStateMeshShardings())
+    self.assertIn("lora_p", prev_params)
+    self.assertNotIn("p", prev_params)
+
+
+class TestRestoreLoraNnx(unittest.TestCase):
+  """Unit tests for lora_utils.restore_lora_from_path under NNX."""
+
+  def test_restore_lora_from_standalone_and_full_trainstate(self):
+    # 1. Setup a toy model with LoRA parameters
+    class ToyModel(nnx.Module):
+
+      def __init__(self):
+        self.p = nnx.Param(jnp.ones((2, 2)))
+        self.lora_p = nnx.LoRAParam(jnp.zeros((2, 2)))
+
+    class DummyLoraConfig:
+      enable_lora = True
+      lora_restore_path = "dummy/restore/path"
+      lora_rank = 4
+      lora_alpha = 8.0
+
+    class DummyConfig:
+      lora = DummyLoraConfig()
+      scan_layers = False
+
+    model = ToyModel()
+
+    # 2. Mocking Orbax restore for standalone LoRA checkpoint
+    standalone_restored_state = {"lora_p": {"value": jnp.ones((2, 2)) * 5.0}}
+
+    with mock.patch(
+        "orbax.checkpoint.PyTreeCheckpointer.restore",
+        return_value=standalone_restored_state,
+    ) as mock_restore:
+      lora_utils.restore_lora_from_path(model, DummyConfig())
+      mock_restore.assert_called_once()
+      # Verify that lora_p value was updated to 5.0
+      np.testing.assert_allclose(np.asarray(model.lora_p[...]), 5.0)
+
+    # Reset lora_p value
+    model.lora_p[...] = jnp.zeros((2, 2))
+
+    # 3. Mocking Orbax restore for full TrainState checkpoint (dict with "model")
+    full_trainstate_dict = {
+        "model": {"lora_p": {"value": jnp.ones((2, 2)) * 10.0}},
+        "optimizer": {},
+    }
+    with mock.patch(
+        "orbax.checkpoint.PyTreeCheckpointer.restore",
+        return_value=full_trainstate_dict,
+    ) as mock_restore:
+      lora_utils.restore_lora_from_path(model, DummyConfig())
+      mock_restore.assert_called_once()
+      # Verify that lora_p value was updated to 10.0
+      np.testing.assert_allclose(np.asarray(model.lora_p[...]), 10.0)
+
+    # Reset lora_p value
+    model.lora_p[...] = jnp.zeros((2, 2))
+
+    # 4. Mocking Orbax restore for full TrainState checkpoint (object with .model attribute)
+    class DummyTrainState:
+
+      def __init__(self):
+        self.model = {"lora_p": {"value": jnp.ones((2, 2)) * 15.0}}
+
+    dummy_state_obj = DummyTrainState()
+    with mock.patch("orbax.checkpoint.PyTreeCheckpointer.restore", return_value=dummy_state_obj) as mock_restore:
+      lora_utils.restore_lora_from_path(model, DummyConfig())
+      mock_restore.assert_called_once()
+      # Verify that lora_p value was updated to 15.0
+      np.testing.assert_allclose(np.asarray(model.lora_p[...]), 15.0)
+
+    # 5. Mocking Orbax restore for checkpoint missing LoRA parameters (raises ValueError)
+    base_only_checkpoint = {
+        "p": {"value": jnp.ones((2, 2)) * 20.0},
+    }
+    with mock.patch("orbax.checkpoint.PyTreeCheckpointer.restore", return_value=base_only_checkpoint) as mock_restore:
+      with self.assertRaises(ValueError) as context:
+        lora_utils.restore_lora_from_path(model, DummyConfig())
+      self.assertIn("No LoRA/adapter parameters were successfully restored", str(context.exception))
+      mock_restore.assert_called_once()
+
+
+class TestLoraResumeAndConversion(unittest.TestCase):
+  """Tests for setup_initial_state on LoRA resume path and upfront conversion."""
+
+  def test_setup_initial_state_lora_resume(self):
+    """Test setup_initial_state on LoRA resume path."""
+    # Setup mock config
+    config = mock.MagicMock()
+    config.pure_nnx = True
+    config.lora = mock.MagicMock()
+    config.lora.enable_lora = True
+    config.load_parameters_path = "mock/base_params_path"
+    config.load_full_state_path = ""
+    config.checkpoint_storage_concurrent_gb = 8
+    config.checkpoint_storage_use_ocdbt = False
+    config.checkpoint_storage_use_zarr3 = False
+
+    # Mock State with model containing nnx.Param and nnx.LoRAParam
+    class MockModel(nnx.Module):
+
+      def __init__(self):
+        self.kernel = nnx.Param(jax.numpy.ones((2, 2)))
+        self.lora_a = nnx.LoRAParam(jax.numpy.zeros((2, 2)))
+
+    class MockTrainState(nnx.Module):
+
+      def __init__(self, model):
+        self.model = model
+
+    mock_model = MockModel()
+    state_mesh_shardings = mock.MagicMock()
+
+    # init_state_fn mock returning a proper nnx.Module container
+    def init_state_fn():
+      return MockTrainState(mock_model)
+
+    # Simulating Orbax restoration behavior: missing base weights are restored as ShapeDtypeStructs
+    restored_pure_dict = nnx.state(MockTrainState(mock_model), nnx.LoRAParam).to_pure_dict()
+    restored_pure_dict["model"] = {
+        "kernel": jax.ShapeDtypeStruct((2, 2), jax.numpy.float32),
+        "lora_a": restored_pure_dict["model"]["lora_a"],
+    }
+    restored = {"items": nnx.State(restored_pure_dict)}
+
+    # Mock checkpointing.load_params_from_path to return updated base params
+    loaded_base = nnx.State({"kernel": nnx.Param(jax.numpy.ones((2, 2)) * 5.0)})
+
+    with (
+        mock.patch(
+            "maxtext.utils.maxtext_utils.get_abstract_state",
+            return_value=(mock.MagicMock(), mock.MagicMock(), state_mesh_shardings),
+        ),
+        mock.patch("maxtext.common.checkpointing.load_state_if_possible", return_value=(restored, None)),
+        mock.patch("maxtext.common.checkpointing.load_params_from_path", return_value=loaded_base) as mock_load,
+        mock.patch("jax.jit", side_effect=lambda f, *args, **kwargs: f),
+    ):
+      state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, config, None, None, init_state_fn, True)
+      # The base model should be restored, meaning its kernel has updated value of 5.0
+      mock_load.assert_called_once()
+      self.assertTrue(was_restored)
+      self.assertEqual(state.model.kernel.get_value()[0, 0], 5.0)
+
+  def test_setup_initial_state_upfront_lora_conversion(self):
+    """Test upfront conversion of quantized base weights on LoRA path."""
+    config = mock.MagicMock()
+    config.pure_nnx = True
+    config.lora = mock.MagicMock()
+    config.lora.enable_lora = True
+    config.load_parameters_path = ""
+    config.load_full_state_path = ""
+    config.checkpoint_storage_concurrent_gb = 8
+    config.checkpoint_storage_use_ocdbt = False
+    config.checkpoint_storage_use_zarr3 = False
+
+    # Mock variables representing quantized weights as nnx.State
+    quantized_state = nnx.State({"qvalue": jax.numpy.ones((2, 2)), "scale": jax.numpy.ones((1, 2))})
+
+    class MockModel(nnx.Module):
+
+      def __init__(self):
+        # Initial value of Param holds the State
+        self.kernel = nnx.Param(quantized_state)
+
+    class MockTrainState(nnx.Module):
+
+      def __init__(self, model):
+        self.model = model
+
+    mock_model = MockModel()
+    state_mesh_shardings = mock.MagicMock()
+
+    # init_state_fn mock returning a proper nnx.Module container
+    def init_state_fn():
+      return MockTrainState(mock_model)
+
+    with (
+        mock.patch(
+            "maxtext.utils.maxtext_utils.get_abstract_state",
+            return_value=(mock.MagicMock(), mock.MagicMock(), state_mesh_shardings),
+        ),
+        mock.patch("maxtext.common.checkpointing.load_state_if_possible", return_value=(None, None)),
+        mock.patch("jax.jit", side_effect=lambda f, *args, **kwargs: f),
+    ):
+      state, _, _, _, _ = maxtext_utils.setup_initial_state(None, config, None, None, init_state_fn, True)
+
+    # Assert kernel value is converted from nnx.State to QArray
+    self.assertTrue(isinstance(state.model.kernel.get_value(), QArray))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/lora_utils_nnx_test.py b/tests/unit/lora_utils_nnx_test.py
deleted file mode 100644
index 322425e674..0000000000
--- a/tests/unit/lora_utils_nnx_test.py
+++ /dev/null
@@ -1,294 +0,0 @@
-# Copyright 2023–2026 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Unit tests for the NNX-shaped LoRA helpers in `lora_utils`, with a small
-Linen regression block at the end."""
-
-import unittest
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from maxtext.utils.lora_utils import (
-    apply_lora_on_base_params,
-    apply_lora_on_base_params_nnx,
-    get_lora_abstract_state_nnx,
-    unapply_lora_from_base_params,
-    unapply_lora_from_base_params_nnx,
-)
-
-
-# ---------------------------------------------------------------------------
-# Fake abstract state builders (mirror the NNX vs. Linen tree shapes)
-# ---------------------------------------------------------------------------
-
-
-def _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4, dtype=jnp.float32):
-  """Tiny NNX-shaped abstract state for one attention block."""
-
-  def _sds(shape):
-    return jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=None)
-
-  return {
-      "decoder": {
-          "layers": {
-              "self_attention": {
-                  "query": {"kernel": _sds((emb, num_heads, head_dim))},
-                  "key": {"kernel": _sds((emb, num_heads, head_dim))},
-                  "value": {"kernel": _sds((emb, num_heads, head_dim))},
-                  "out": {"kernel": _sds((emb, num_heads, head_dim))},
-              },
-              "mlp": {"wi": {"kernel": _sds((emb, 4 * emb))}},
-          },
-          "shared_embedding": {"embedding": _sds((100, emb))},
-      },
-  }
-
-
-def _make_linen_attention_abstract(emb=8, num_heads=2, head_dim=4, dtype=jnp.float32):
-  """Linen-shaped equivalent (with the `{"params": ...}` outer wrap)."""
-  return {"params": _make_nnx_attention_abstract(emb, num_heads, head_dim, dtype)}
-
-
-def _lora_config(rank=4, alpha=8.0, target_modules=("q_proj", "v_proj")):
-  return {
-      "r": rank,
-      "lora_alpha": alpha,
-      "target_modules": list(target_modules),
-  }
-
-
-# ---------------------------------------------------------------------------
-# get_lora_abstract_state_nnx
-# ---------------------------------------------------------------------------
-
-
-class TestGetLoraAbstractStateNnx(unittest.TestCase):
-  """`get_lora_abstract_state_nnx` shape, sharding, and error-path coverage."""
-
-  def test_lora_shapes_for_query_and_value(self):
-    abs_params = _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4)
-    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(rank=4))
-    attn = state.params["decoder"]["layers"]["self_attention"]
-
-    a = attn["query"]["lora_a.kernel"]
-    b = attn["query"]["lora_b.kernel"]
-    self.assertEqual(a.shape, (8, 4))
-    self.assertEqual(b.shape, (4, 2, 4))
-    self.assertEqual(a.dtype, jnp.float32)
-    self.assertEqual(b.dtype, jnp.float32)
-
-    a = attn["value"]["lora_a.kernel"]
-    b = attn["value"]["lora_b.kernel"]
-    self.assertEqual(a.shape, (8, 4))
-    self.assertEqual(b.shape, (4, 2, 4))
-
-  def test_non_target_modules_emit_none_leaves(self):
-    abs_params = _make_nnx_attention_abstract()
-    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(target_modules=("q_proj",)))
-    attn = state.params["decoder"]["layers"]["self_attention"]
-    self.assertIn("lora_a.kernel", attn["query"])
-    self.assertIsNone(attn["key"]["kernel"])
-    self.assertIsNone(attn["value"]["kernel"])
-    self.assertIsNone(attn["out"]["kernel"])
-    self.assertIsNone(state.params["decoder"]["layers"]["mlp"]["wi"]["kernel"])
-    self.assertIsNone(state.params["decoder"]["shared_embedding"]["embedding"])
-
-  def test_o_proj_has_distinct_shape(self):
-    abs_params = _make_nnx_attention_abstract(emb=8, num_heads=2, head_dim=4)
-    state, _ = get_lora_abstract_state_nnx(abs_params, _lora_config(rank=3, target_modules=("o_proj",)))
-    out = state.params["decoder"]["layers"]["self_attention"]["out"]
-    a = out["lora_a.kernel"]
-    b = out["lora_b.kernel"]
-    # For a 3D base (emb, num_heads, head_dim): lora_a.shape ends with rank,
-    # lora_b shape is (rank, head_dim).
-    self.assertEqual(a.shape, (8, 2, 3))
-    self.assertEqual(b.shape, (3, 4))
-
-  def test_unsupported_leaf_type_raises(self):
-    bad = {"decoder": {"layers": {"self_attention": {"query": {"kernel": jnp.zeros((4, 2, 2))}}}}}
-    with self.assertRaises(ValueError):
-      get_lora_abstract_state_nnx(bad, _lora_config())
-
-  def test_unexpected_leaf_name_raises(self):
-    bad = {"decoder": {"layers": {"self_attention": {"query": {"weight": jax.ShapeDtypeStruct((4, 2), jnp.float32)}}}}}
-    with self.assertRaises(ValueError):
-      get_lora_abstract_state_nnx(bad, _lora_config())
-
-  # Linen-vs-NNX numerical parity is covered by TestApplyLoraNnx.test_numerical_parity_with_linen_apply.
-
-
-# ---------------------------------------------------------------------------
-# apply / unapply on NNX-shape pure dicts
-# ---------------------------------------------------------------------------
-
-
-def _concrete_base(rng=None, emb=4, num_heads=2, head_dim=3):
-  """Concrete arrays mirroring the abstract structure used above (NNX-shape)."""
-  if rng is None:
-    rng = jax.random.key(0)
-  k1, k2, k3, k4, k5, k6 = jax.random.split(rng, 6)
-  shape_attn = (emb, num_heads, head_dim)
-  return {
-      "decoder": {
-          "layers": {
-              "self_attention": {
-                  "query": {"kernel": jax.random.normal(k1, shape_attn)},
-                  "key": {"kernel": jax.random.normal(k2, shape_attn)},
-                  "value": {"kernel": jax.random.normal(k3, shape_attn)},
-                  "out": {"kernel": jax.random.normal(k4, shape_attn)},
-              },
-              "mlp": {"wi": {"kernel": jax.random.normal(k5, (emb, 4 * emb))}},
-          },
-          "shared_embedding": {"embedding": jax.random.normal(k6, (100, emb))},
-      },
-  }
-
-
-def _build_lora_params(base, lora_config_dict, rng):
-  """Build a concrete LoRA tree (random arrays) matching `base`."""
-  abs_tree = jax.tree_util.tree_map(lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=None), base)
-  lora_state, _ = get_lora_abstract_state_nnx(abs_tree, lora_config_dict)
-
-  def _to_concrete(leaf, rng_key):
-    if leaf is None:
-      return None
-    return jax.random.normal(rng_key, leaf.shape, leaf.dtype)
-
-  leaves, tree = jax.tree_util.tree_flatten(lora_state.params, is_leaf=lambda x: x is None)
-  rngs = jax.random.split(rng, max(1, len(leaves)))
-  out_leaves = [_to_concrete(l, r) for l, r in zip(leaves, rngs)]
-  return jax.tree_util.tree_unflatten(tree, out_leaves)
-
-
-class TestApplyLoraNnx(unittest.TestCase):
-  """`apply_lora_on_base_params_nnx` round-trip and Linen-vs-NNX parity."""
-
-  def test_apply_then_unapply_is_identity(self):
-    rng = jax.random.key(42)
-    base_orig = _concrete_base(rng)
-    base = jax.tree_util.tree_map(jnp.copy, base_orig)
-    lora = _build_lora_params(base, _lora_config(rank=2, target_modules=("q_proj", "v_proj")), jax.random.key(7))
-    apply_lora_on_base_params_nnx(base, lora, lora_scale_factor=0.5)
-    # The query and value kernels are targets and must have changed.
-    self.assertFalse(
-        jnp.allclose(
-            base["decoder"]["layers"]["self_attention"]["query"]["kernel"],
-            base_orig["decoder"]["layers"]["self_attention"]["query"]["kernel"],
-        )
-    )
-    # The key and out kernels are non-targets and must be untouched.
-    np.testing.assert_array_equal(
-        np.asarray(base["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
-        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
-    )
-    np.testing.assert_array_equal(
-        np.asarray(base["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
-        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
-    )
-    unapply_lora_from_base_params_nnx(base, lora, lora_scale_factor=0.5)
-    np.testing.assert_allclose(
-        np.asarray(base["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        rtol=1e-5,
-        atol=1e-6,
-    )
-    np.testing.assert_allclose(
-        np.asarray(base["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
-        np.asarray(base_orig["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
-        rtol=1e-5,
-        atol=1e-6,
-    )
-
-  def test_numerical_parity_with_linen_apply(self):
-    """The NNX and Linen apply paths produce identical results on the same inputs."""
-    rng = jax.random.key(123)
-    base_nnx = _concrete_base(rng)
-    base_linen = {"params": jax.tree_util.tree_map(jnp.copy, base_nnx)}
-    lora = _build_lora_params(base_nnx, _lora_config(rank=2, target_modules=("q_proj",)), jax.random.key(5))
-    apply_lora_on_base_params_nnx(base_nnx, lora, lora_scale_factor=0.7)
-    apply_lora_on_base_params(base_linen, {"params": lora}, lora_scale_factor=0.7)
-    np.testing.assert_allclose(
-        np.asarray(base_nnx["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        np.asarray(base_linen["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        rtol=1e-6,
-    )
-
-  def test_apply_with_unexpected_lora_key_raises(self):
-    base = _concrete_base()
-    bad = {"decoder": {"layers": {"self_attention": {"query": {"unexpected": jnp.zeros((4, 2))}}}}}
-    with self.assertRaises(ValueError):
-      apply_lora_on_base_params_nnx(base, bad)
-
-
-class TestLinenLoraRegression(unittest.TestCase):
-  """Smoke tests for the Linen apply / unapply helpers (no other unit test exercises them)."""
-
-  def _linen_pair(self, rng=None):
-    """Build a Linen-shape (with `{"params": ...}` outer wrapper) base + lora pair."""
-    if rng is None:
-      rng = jax.random.key(99)
-    base_inner = _concrete_base(rng)
-    base = {"params": jax.tree_util.tree_map(jnp.copy, base_inner)}
-    lora_inner = _build_lora_params(
-        base_inner,
-        _lora_config(rank=2, target_modules=("q_proj", "v_proj")),
-        jax.random.key(7),
-    )
-    lora = {"params": lora_inner}
-    return base, lora
-
-  def test_linen_apply_then_unapply_is_identity(self):
-    base, lora = self._linen_pair()
-    base_orig = jax.tree_util.tree_map(jnp.copy, base)
-    apply_lora_on_base_params(base, lora, lora_scale_factor=0.5)
-    unapply_lora_from_base_params(base, lora, lora_scale_factor=0.5)
-    np.testing.assert_allclose(
-        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"]),
-        rtol=1e-5,
-        atol=1e-6,
-    )
-    np.testing.assert_allclose(
-        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
-        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["value"]["kernel"]),
-        rtol=1e-5,
-        atol=1e-6,
-    )
-
-  def test_linen_apply_only_modifies_target_modules(self):
-    base, lora = self._linen_pair()
-    base_orig = jax.tree_util.tree_map(jnp.copy, base)
-    apply_lora_on_base_params(base, lora, lora_scale_factor=1.0)
-    # query and value are targets and must change.
-    self.assertFalse(
-        jnp.allclose(
-            base["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"],
-            base_orig["params"]["decoder"]["layers"]["self_attention"]["query"]["kernel"],
-        )
-    )
-    # key and out are non-target and must be untouched.
-    np.testing.assert_array_equal(
-        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
-        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["key"]["kernel"]),
-    )
-    np.testing.assert_array_equal(
-        np.asarray(base["params"]["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
-        np.asarray(base_orig["params"]["decoder"]["layers"]["self_attention"]["out"]["kernel"]),
-    )
-
-
-if __name__ == "__main__":
-  unittest.main()
diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py
index ae33d66d61..0c6e3c138a 100644
--- a/tests/unit/train_state_nnx_checkpoint_test.py
+++ b/tests/unit/train_state_nnx_checkpoint_test.py
@@ -363,6 +363,7 @@ def _config(self, **overrides):
         "enable_multi_tier_checkpointing": False,
         "local_checkpoint_period": 0,
         "enable_autocheckpoint": False,
+        "lora": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -505,16 +506,13 @@ def test_maybe_save_checkpoint_skips_non_checkpoint_step_before_state_work(
     mgr = mock.MagicMock()
     mgr.reached_preemption.return_value = False
 
-    with (
-        mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock,
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-    ):
+    with mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock:
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=3)
 
     mgr.latest_step.assert_not_called()
     mgr.reached_preemption.assert_called_once_with(3)
     mgr.wait_until_finished.assert_not_called()
-    to_checkpoint_dict_mock.assert_not_called()
+    state.to_pure_dict.assert_not_called()
     save_checkpoint_mock.assert_not_called()
 
   def test_maybe_save_checkpoint_handles_preemption_on_non_checkpoint_step(
@@ -526,17 +524,14 @@ def test_maybe_save_checkpoint_handles_preemption_on_non_checkpoint_step(
     mgr = mock.MagicMock()
     mgr.reached_preemption.return_value = True
 
-    with (
-        mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock,
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-    ):
+    with mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock:
       with self.assertRaises(checkpointing.exceptions.StopTraining):
         checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=3)
 
     mgr.latest_step.assert_not_called()
     mgr.reached_preemption.assert_called_once_with(3)
     mgr.wait_until_finished.assert_called_once_with()
-    to_checkpoint_dict_mock.assert_not_called()
+    state.to_pure_dict.assert_not_called()
     save_checkpoint_mock.assert_not_called()
 
   def test_maybe_save_checkpoint_allows_local_checkpoint_period(self):
@@ -546,7 +541,7 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self):
         "enable_multi_tier_checkpointing",
     ):
       with self.subTest(checkpoint_flag=checkpoint_flag):
-        state = mock.Mock()
+        state = self._build_nnx_state(5)
         config = self._config(
             checkpoint_period=100,
             local_checkpoint_period=5,
@@ -557,23 +552,25 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self):
         mgr.reached_preemption.return_value = False
         save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-        with (
-            mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-            mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-        ):
+        with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
           checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
         mgr.latest_step.assert_called_once_with()
         mgr.reached_preemption.assert_called_once_with(5)
         mgr.wait_until_finished.assert_not_called()
-        to_checkpoint_dict_mock.assert_called_once_with(state)
         save_checkpoint_mock.assert_called_once()
 
+        # Verify the actual formatted checkpoint dictionary is passed to save_checkpoint
+        saved_dict = save_checkpoint_mock.call_args[0][2]
+        self.assertIn("params", saved_dict)
+        self.assertIn("step", saved_dict)
+        self.assertIn("opt_state", saved_dict)
+
   def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
       self,
   ):
     """Continuous checkpointing should not suppress MTC local saves."""
-    state = mock.Mock()
+    state = self._build_nnx_state(5)
     config = self._config(
         checkpoint_period=100,
         enable_continuous_checkpointing=True,
@@ -586,18 +583,20 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
     mgr.reached_preemption.return_value = False
     save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-    with (
-        mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-    ):
+    with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
     mgr.should_save.assert_called_once_with(5)
     mgr.latest_step.assert_called_once_with()
     mgr.reached_preemption.assert_called_once_with(5)
-    to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
 
+    # Verify the actual formatted checkpoint dictionary is passed to save_checkpoint
+    saved_dict = save_checkpoint_mock.call_args[0][2]
+    self.assertIn("params", saved_dict)
+    self.assertIn("step", saved_dict)
+    self.assertIn("opt_state", saved_dict)
+
 
 class TestLinenCheckpointFormatConverters(unittest.TestCase):
   """to_linen_checkpoint_dict / from_linen_checkpoint_dict (NNX <-> Linen on-disk layout)."""