From 3b99bb5d9d2688379b9c080668a01cbc5af81d0a Mon Sep 17 00:00:00 2001
From: Jacky Fang <rexbear87941@gmail.com>
Date: Fri, 3 Jul 2026 03:43:06 +0000
Subject: [PATCH 1/4] feat(nnx): support native Flax NNX PEFT/LoRA training
 loop

---
 src/maxtext/utils/lora_utils.py               |  9 +-
 src/maxtext/utils/sharding.py                 | 14 +--
 src/maxtext/utils/train_utils.py              | 13 ++-
 .../integration/setup_train_loop_nnx_test.py  | 97 +++++++++++++++++--
 tests/unit/lora_utils_nnx_test.py             | 30 ++++++
 5 files changed, 149 insertions(+), 14 deletions(-)

diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 1aa443f093..14d077a7ac 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -607,7 +607,14 @@ def apply_lora_to_model(
   )
 
   if mesh is not None:
-    with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+    from contextlib import nullcontext
+    try:
+      from jax._src import core as jax_core
+      mesh_ctx = jax.set_mesh(mesh) if jax_core.trace_state_clean() else nullcontext()
+    except Exception:
+      mesh_ctx = nullcontext()
+
+    with mesh_ctx, 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.
diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py
index dc2f70ae14..bf6c633bc8 100644
--- a/src/maxtext/utils/sharding.py
+++ b/src/maxtext/utils/sharding.py
@@ -609,19 +609,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 (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:
@@ -629,7 +631,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 73e5d06b05..3b18bddad0 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
 
 
@@ -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
@@ -367,6 +372,12 @@ def create_train_state_fn():
     else:
       train_state = nnx.merge(state_graphdef, state)
       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
 
diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py
index f446c8f9df..f5a7803249 100644
--- a/tests/integration/setup_train_loop_nnx_test.py
+++ b/tests/integration/setup_train_loop_nnx_test.py
@@ -44,6 +44,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,14 +123,98 @@ 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.assertGreater(len(jax.tree.leaves(params)), 0)
-
     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",
+        lora={
+            "enable_lora": True,
+            "lora_rank": 4,
+            "lora_alpha": 8.0,
+            "lora_module_path": ".*decoder.*",
+        }
+    )
+    (
+        _,
+        _,
+        _,
+        model,
+        _,
+        _,
+        _,
+        _,
+        _,
+        _,
+        train_state,
+    ) = setup_train_loop(config, recorder=None)
+    
+    from maxtext.utils.lora_utils import is_lora_enabled
+    self.assertTrue(is_lora_enabled(model))
+    self.assertIs(train_state.optimizer.wrt, nnx.LoRAParam)
+
+  def test_train_step_updates_only_lora_weights(self):
+    import jax.numpy as jnp
+    config = _tiny_nnx_pyconfig(
+        weight_dtype="bfloat16",
+        lora={
+            "enable_lora": True,
+            "lora_rank": 4,
+            "lora_alpha": 8.0,
+            "lora_module_path": ".*decoder.*",
+        },
+        steps=1,
+    )
+    (
+        _,
+        _,
+        _,
+        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
+    from maxtext.trainers.pre_train import train
+    batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
+    
+    new_state, metrics = 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))
+
+
 if __name__ == "__main__":
   unittest.main()
+
diff --git a/tests/unit/lora_utils_nnx_test.py b/tests/unit/lora_utils_nnx_test.py
index 322425e674..f3a491d1fd 100644
--- a/tests/unit/lora_utils_nnx_test.py
+++ b/tests/unit/lora_utils_nnx_test.py
@@ -290,5 +290,35 @@ def test_linen_apply_only_modifies_target_modules(self):
     )
 
 
+class TestShardingExtractionNnx(unittest.TestCase):
+  """Test sharding parameter extraction under different LoRA configurations."""
+
+  def test_sharding_extracts_only_lora_params(self):
+    from maxtext.utils import sharding
+    from flax import nnx
+    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()
+    graphdef, 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)
+
+
 if __name__ == "__main__":
   unittest.main()
+

From aa01607817bb60bc409a227aa39d71b491753041 Mon Sep 17 00:00:00 2001
From: Jacky Fang <rexbear87941@gmail.com>
Date: Fri, 3 Jul 2026 03:55:43 +0000
Subject: [PATCH 2/4] refactor(lora): remove contextlib dependency and fix
 sharding tolerance for integration tests

---
 src/maxtext/utils/lora_utils.py                | 13 ++++++++++---
 tests/integration/setup_train_loop_nnx_test.py |  2 ++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 14d077a7ac..9a639e7539 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -607,12 +607,19 @@ def apply_lora_to_model(
   )
 
   if mesh is not None:
-    from contextlib import nullcontext
     try:
       from jax._src import core as jax_core
-      mesh_ctx = jax.set_mesh(mesh) if jax_core.trace_state_clean() else nullcontext()
+      is_trace_clean = jax_core.trace_state_clean()
     except Exception:
-      mesh_ctx = nullcontext()
+      is_trace_clean = False
+
+    class _NullContext:
+      def __enter__(self):
+        return None
+      def __exit__(self, *args):
+        pass
+
+    mesh_ctx = jax.set_mesh(mesh) if is_trace_clean else _NullContext()
 
     with mesh_ctx, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
       graph_def, state = nnx.split(lora_model)
diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py
index f5a7803249..7c7cf38735 100644
--- a/tests/integration/setup_train_loop_nnx_test.py
+++ b/tests/integration/setup_train_loop_nnx_test.py
@@ -133,6 +133,7 @@ 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,
@@ -162,6 +163,7 @@ def test_train_step_updates_only_lora_weights(self):
     import jax.numpy as jnp
     config = _tiny_nnx_pyconfig(
         weight_dtype="bfloat16",
+        sharding_tolerance=1.0,
         lora={
             "enable_lora": True,
             "lora_rank": 4,

From 2e4371caf062f90748d037b80fe1392a35e4ed33 Mon Sep 17 00:00:00 2001
From: Jacky Fang <rexbear87941@gmail.com>
Date: Fri, 3 Jul 2026 04:14:26 +0000
Subject: [PATCH 3/4] refactor(lora): remove redundant jax.set_mesh context
 manager in apply_lora_to_model

---
 src/maxtext/utils/lora_utils.py | 16 +---------------
 1 file changed, 1 insertion(+), 15 deletions(-)

diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 9a639e7539..e9806d6cb7 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -607,21 +607,7 @@ def apply_lora_to_model(
   )
 
   if mesh is not None:
-    try:
-      from jax._src import core as jax_core
-      is_trace_clean = jax_core.trace_state_clean()
-    except Exception:
-      is_trace_clean = False
-
-    class _NullContext:
-      def __enter__(self):
-        return None
-      def __exit__(self, *args):
-        pass
-
-    mesh_ctx = jax.set_mesh(mesh) if is_trace_clean else _NullContext()
-
-    with mesh_ctx, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+    with nn_partitioning.axis_rules(mt_config.logical_axis_rules):
       graph_def, state = nnx.split(lora_model)
 
       # We handle explicit replication for LoRA to ensure safety and efficiency.

From 4fe1fc83d3ca639ce158f9a40c422006df883763 Mon Sep 17 00:00:00 2001
From: Jacky Fang <jacky.fang@cienet.com>
Date: Mon, 6 Jul 2026 09:34:55 +0000
Subject: [PATCH 4/4] test(nnx): fix LoRA integration test by removing
 lora_module_path override

---
 src/maxtext/layers/linears.py                 |  42 +++++++-
 src/maxtext/trainers/pre_train/train.py       |  30 +++---
 src/maxtext/utils/lora_utils.py               |  52 +++++++--
 src/maxtext/utils/sharding.py                 |  30 ++----
 src/maxtext/utils/train_utils.py              |  14 +--
 .../integration/setup_train_loop_nnx_test.py  |  61 +++++++----
 tests/unit/lora_utils_nnx_test.py             | 101 ++++++++++++++++--
 7 files changed, 242 insertions(+), 88 deletions(-)

diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py
index ca99776f38..180d1f13ce 100644
--- a/src/maxtext/layers/linears.py
+++ b/src/maxtext/layers/linears.py
@@ -70,6 +70,35 @@ def canonicalize_tuple(x):
     return (x,)
 
 
+def _restore_custom_types_from_state(val):
+  """Restores qwix custom quantized types from nnx.State representation."""
+  if isinstance(val, nnx.State):
+    pure_dict = {k: _restore_custom_types_from_state(v) for k, v in val.items()}
+    if "array" in pure_dict:
+      # pylint: disable=import-outside-toplevel
+      from qwix._src.providers.ptq import WithAux
+
+      return WithAux(array=pure_dict["array"], how=None)
+    elif "qvalue" in pure_dict and "scale" in pure_dict:
+      # pylint: disable=import-outside-toplevel
+      from qwix._src.core.qarray import QArray
+
+      return QArray(
+          qvalue=pure_dict["qvalue"],
+          scale=pure_dict["scale"],
+          zero_point=pure_dict.get("zero_point", None),
+          qtype=None,
+      )
+    else:
+      return pure_dict
+  elif isinstance(val, nnx.Variable):
+    return _restore_custom_types_from_state(val.get_value())
+  elif isinstance(val, dict):
+    return {k: _restore_custom_types_from_state(v) for k, v in val.items()}
+  else:
+    return val
+
+
 def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul_precision, quant):
   """Computes a dot_general operation that may be quantized."""
   dot_general = lax.dot_general
@@ -222,8 +251,16 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
     if quantizations.in_serve_mode(self.quant):
       kernel_shape = self.in_features_shape + self.out_features_shape
       kernel = jnp.zeros(kernel_shape, dtype=self.dtype)
+      restored_to_state = False
     else:
-      kernel = self.kernel[...]
+      kernel_val = self.kernel._raw_value if hasattr(self.kernel, "_raw_value") else self.kernel
+      restored_to_state = False
+      if isinstance(kernel_val, nnx.State):
+        kernel = _restore_custom_types_from_state(kernel_val)
+        self.kernel.value = kernel
+        restored_to_state = True
+      else:
+        kernel = self.kernel[...]
       # Move logit_dense kernel to device if parameter offloading is enabled
       if self.parameter_memory_host_offload:
         max_logging.log("linear.py: Moving parameter logits_dense kernel to device")
@@ -246,6 +283,9 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
         out_sharding,
     )
 
+    if restored_to_state:
+      self.kernel.value = kernel_val
+
     if self.bias is not None:
       bias = jnp.asarray(self.bias[...], self.dtype)
       output += bias
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e7f5ea84b6..94b2b9a92f 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -78,9 +78,7 @@
 def get_first_step(model, state):
   if isinstance(model, nn.Module):
     return int(state.step)
-  if hasattr(
-      state, "inner_state"
-  ):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
+  if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
     return int(state.step.get_value())
   return int(state.optimizer.step.get_value())
 
@@ -393,9 +391,10 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
           is_train=True,
       )
     else:
+      wrt = nnx.LoRAParam if config.lora.enable_lora else nnx.Param
       owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True)
       custom_param_filter = nnx.Any(owg_type)
-      model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, nnx.Param, custom_param_filter, ...)
+      model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, wrt, custom_param_filter, ...)
       if config.parameter_memory_host_offload:
         # Params are kept on host (pinned_host) in in_shardings. Move only Param
         # variables to device before the forward/backward pass so that all dot_general
@@ -420,7 +419,7 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
       def diff_wrapper(curr_params, custom_params, rest, config, data):
         local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
         loss, aux = loss_fn(local_model, config, data, None, None, is_train=True)
-        _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...)
+        _, _, _, new_rest = nnx.split(local_model, wrt, custom_param_filter, ...)
         return loss, (aux, new_rest)
 
       grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True)
@@ -771,6 +770,17 @@ def train_loop(config, recorder, state=None):
   start_step = get_first_step(model, state)  # this is the start_step for training
   train_utils.validate_completed_steps(start_step, config.steps)
 
+  if config.pure_nnx:
+    if isinstance(state, nnx.Module):
+      _ = nnx.pop(state, nnx.Intermediate)
+    elif isinstance(state, nnx.State):
+      state = nnx.state(state, nnx.Not(nnx.Intermediate))
+
+    if isinstance(state_mesh_shardings, nnx.Module):
+      _ = nnx.pop(state_mesh_shardings, nnx.Intermediate)
+    elif isinstance(state_mesh_shardings, nnx.State):
+      state_mesh_shardings = nnx.state(state_mesh_shardings, nnx.Not(nnx.Intermediate))
+
   if isinstance(model, nn.Module):
     jit_model = model
   elif config.enable_diloco:
@@ -784,11 +794,7 @@ def train_loop(config, recorder, state=None):
     # the Zero-1 opt overlay doesn't apply through the diloco wrapper.
     params_shardings = state_mesh_shardings.params
   else:
-    params_shardings, state_mesh_shardings = (
-        sharding.maybe_update_params_sharding_with_opt(
-            config, state_mesh_shardings
-        )
-    )
+    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config,
@@ -828,9 +834,7 @@ def train_loop(config, recorder, state=None):
   if isinstance(model, nn.Module):
     setup_params = state.params
   elif config.enable_diloco:
-    setup_params = (
-        state.params
-    )  # DiLoCoTrainState.params: the outer (global) params
+    setup_params = state.params  # DiLoCoTrainState.params: the outer (global) params
   else:
     _, setup_params, _ = nnx.split(state.model, nnx.Param, ...)
   metric_logger_instance.write_setup_info_to_tensorboard(setup_params)
diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index e9806d6cb7..3d177f58ab 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -424,14 +424,20 @@ def _get_lora_module_path(mt_config: pyconfig.HyperParameters) -> str:
   model_name = mt_config.model_name.lower()
 
   # Find the first matching architecture prefix or use 'default'
-  matched_key = next((k for k in lora_configs if k != "default" and model_name.startswith(k)), "default")
+  matched_key = next(
+      (k for k in lora_configs if k != "default" and model_name.startswith(k)),
+      "default",
+  )
 
   if matched_key == "default":
     max_logging.log(f"Warning: Model '{model_name}' is unverified; falling back to default LoRA path.")
   else:
     max_logging.log(f"Auto-detected lora_module_path for model '{model_name}' (matched: '{matched_key}')")
 
-  raw_path = lora_configs.get(matched_key, "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))")
+  raw_path = lora_configs.get(
+      matched_key,
+      "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))",
+  )
 
   # This regex makes the layer index optional, matching both scanned and unscanned layer paths
   # (e.g. 'layers/0/mlp/...' vs 'layers/mlp/...').
@@ -612,7 +618,11 @@ def apply_lora_to_model(
 
       # We handle explicit replication for LoRA to ensure safety and efficiency.
       state = jax.tree_util.tree_map(
-          lambda x: x.replace(sharding=jax.sharding.PartitionSpec(), out_sharding=None, sharding_names=None)
+          lambda x: x.replace(
+              sharding=jax.sharding.PartitionSpec(),
+              out_sharding=None,
+              sharding_names=None,
+          )
           if isinstance(x, nnx.LoRAParam)
           else x,
           state,
@@ -624,18 +634,24 @@ def apply_lora_to_model(
       dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, mt_config.logical_axis_rules)
 
       def _safe_reshard(var, sharding_spec):
-        if not isinstance(var, nnx.Variable) or not isinstance(sharding_spec, jax.sharding.Sharding):
+        physical_sharding = sharding_spec.get_value() if isinstance(sharding_spec, nnx.Variable) else sharding_spec
+        if not isinstance(var, nnx.Variable) or not isinstance(physical_sharding, jax.sharding.Sharding):
           return var
         val = var.get_value()
-        if not isinstance(val, jax.Array):
+        if not isinstance(val, jax.Array) or isinstance(val, jax.core.Tracer):
           return var
         # make_array_from_callback natively constructs a globally sharded array
         # from the local host arrays, bypassing backend-specific device_put issues
         # on both Pathways and McJAX.
-        resharded_val = jax.make_array_from_callback(val.shape, sharding_spec, lambda idx: val[idx])
+        resharded_val = jax.make_array_from_callback(val.shape, physical_sharding, lambda idx: val[idx])
         return var.replace(value=resharded_val)
 
-      state = jax.tree_util.tree_map(_safe_reshard, state, dst_shardings, is_leaf=lambda x: isinstance(x, nnx.Variable))
+      state = jax.tree_util.tree_map(
+          _safe_reshard,
+          state,
+          dst_shardings,
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
+      )
 
       lora_model = nnx.merge(graph_def, state)
 
@@ -677,7 +693,7 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter
   abstract_lora_params = nnx.state(model, nnx.LoRAParam)
 
   target_for_restore = jax.tree.map(
-      lambda v: {"value": v.value},
+      lambda v: {"value": v.get_value()},
       abstract_lora_params,
       is_leaf=lambda n: isinstance(n, nnx.Variable),
   )
@@ -699,6 +715,12 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter
     max_logging.log(f"Guided restore failed: {e}. Falling back to basic restore.")
     restored_lora_params = ocp.PyTreeCheckpointer().restore(lora_restore_path)
 
+  # If restoring from a full TrainState checkpoint, navigate into the model sub-tree
+  if isinstance(restored_lora_params, dict) and "model" in restored_lora_params:
+    restored_lora_params = restored_lora_params["model"]
+  elif hasattr(restored_lora_params, "model"):
+    restored_lora_params = getattr(restored_lora_params, "model")
+
   # Post processing
   def _map_to_state(path, variable):
     if not isinstance(variable, nnx.Variable):
@@ -722,7 +744,7 @@ def _map_to_state(path, variable):
     else:
       matched_val = curr
 
-    variable.value = matched_val
+    variable[...] = matched_val
 
   jax.tree_util.tree_map_with_path(
       _map_to_state,
@@ -839,7 +861,11 @@ def get_lora_abstract_state_nnx(base_abstract_params, lora_config):
   def get_lora_param_shape(base_array_shape, lora_module):
     if len(base_array_shape) > 4:
       raise ValueError(f"Unsupported base array shape {base_array_shape} (>4D)")
-    if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"):
+    if lora_module in (
+        "self_attention.query",
+        "self_attention.key",
+        "self_attention.value",
+    ):
       lora_a_shape = base_array_shape[:-2] + (lora_rank,)
       lora_b_shape = (lora_rank,) + base_array_shape[1:]
     elif lora_module == "self_attention.out":
@@ -858,7 +884,11 @@ def get_lora_param_sharding(base_param_sharding, lora_module):
     base_pspec = base_param_sharding.spec
     if len(base_pspec) > 4:
       raise ValueError("PartitionSpec size > 4 not supported")
-    if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"):
+    if lora_module in (
+        "self_attention.query",
+        "self_attention.key",
+        "self_attention.value",
+    ):
       lora_a_pspec = jax.sharding.PartitionSpec(*(base_pspec[:-2] + ((),)))
       lora_b_pspec = jax.sharding.PartitionSpec(*(((),) + base_pspec[1:]))
     elif lora_module == "self_attention.out":
diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py
index bf6c633bc8..5fc682436b 100644
--- a/src/maxtext/utils/sharding.py
+++ b/src/maxtext/utils/sharding.py
@@ -176,9 +176,7 @@ def remove_size_one_mesh_axis(spec, mesh):
   return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced)
 
 
-def get_nnx_var_named_sharding_with_scan_axis(
-    v: nnx.Variable, mesh
-) -> nnx.Variable:
+def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable:
   """Compute NamedSharding for an NNX variable, correctly handling the scan axis."""
   val = v.get_value()
   if not hasattr(val, "shape"):
@@ -190,11 +188,7 @@ def get_nnx_var_named_sharding_with_scan_axis(
       return v.replace(jax.tree.map(lambda _: replicated, val))
     return v
   metadata = v.get_metadata()
-  out_sharding = (
-      metadata.get("out_sharding")
-      or metadata.get("sharding_names")
-      or metadata.get("sharding")
-  )
+  out_sharding = metadata.get("out_sharding") or metadata.get("sharding_names") or metadata.get("sharding")
   if not out_sharding:
     pspec = P()
   else:
@@ -202,11 +196,7 @@ def get_nnx_var_named_sharding_with_scan_axis(
     if nnx.PARTITION_NAME in metadata:
       partition_name = metadata[nnx.PARTITION_NAME]
       scan_axis = metadata.get("param_scan_axis", 0)
-      out_sharding = (
-          [out_sharding]
-          if isinstance(out_sharding, str)
-          else list(out_sharding)
-      )
+      out_sharding = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding)
       if partition_name not in out_sharding:
         out_sharding.insert(scan_axis, partition_name)
       out_sharding = tuple(out_sharding)
@@ -215,9 +205,7 @@ def get_nnx_var_named_sharding_with_scan_axis(
     local_rules = metadata.get("sharding_rules", ())
     if context_rules or local_rules:
       local_rules_list = list(local_rules) if local_rules is not None else []
-      context_rules_list = (
-          list(context_rules) if context_rules is not None else []
-      )
+      context_rules_list = list(context_rules) if context_rules is not None else []
       rules = local_rules_list + context_rules_list
       pspec = logical_to_mesh_axes(out_sharding, mesh, rules=rules)
     else:
@@ -609,7 +597,11 @@ 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
+  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 trainable parameters from an nnx.State into a nested plain dict.
@@ -703,9 +695,7 @@ def _update_model_var(path, var):
   return prev_params_shardings, updated_state
 
 
-def build_zero1_input_state_mesh_shardings(
-    config, state_mesh_shardings, params_shardings
-):
+def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_shardings):
   """Build the train-step input shardings under shard_optimizer_over_data (Zero-1).
 
   Model params on input use the original pre-Zero-1 sharding (params_shardings),
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 3b18bddad0..7ee0f77feb 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/src/maxtext/utils/train_utils.py
@@ -326,23 +326,15 @@ def create_train_state_fn():
         state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh)
 
         # create state_mesh_shardings for the DilocoState
-        step_mesh = (
-            state_mesh_shardings.optimizer.step.mesh
-            if config.pure_nnx
-            else state_mesh_shardings.step.mesh
-        )
+        step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh
         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()
-            if config.pure_nnx
-            else state_mesh_shardings_params,
+            state_mesh_shardings_params.to_pure_dict() if config.pure_nnx else state_mesh_shardings_params,
             outer_opt_state_sharding,
-            jax.sharding.NamedSharding(
-                mesh=step_mesh, spec=jax.sharding.PartitionSpec()
-            ),
+            jax.sharding.NamedSharding(mesh=step_mesh, spec=jax.sharding.PartitionSpec()),
         )
 
     # TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal
diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py
index 7c7cf38735..7e9d1f5cfb 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):
@@ -123,6 +125,7 @@ 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(params.keys(), params_shardings.keys())
     del model
 
 
@@ -138,8 +141,7 @@ def test_pure_nnx_setup_with_lora_enabled(self):
             "enable_lora": True,
             "lora_rank": 4,
             "lora_alpha": 8.0,
-            "lora_module_path": ".*decoder.*",
-        }
+        },
     )
     (
         _,
@@ -154,23 +156,32 @@ def test_pure_nnx_setup_with_lora_enabled(self):
         _,
         train_state,
     ) = setup_train_loop(config, recorder=None)
-    
-    from maxtext.utils.lora_utils import is_lora_enabled
+
     self.assertTrue(is_lora_enabled(model))
     self.assertIs(train_state.optimizer.wrt, nnx.LoRAParam)
 
-  def test_train_step_updates_only_lora_weights(self):
-    import jax.numpy as jnp
+  def _run_train_step_updates_only_adapters_test(self, qlora: 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,
+          }
+      )
+
     config = _tiny_nnx_pyconfig(
         weight_dtype="bfloat16",
         sharding_tolerance=1.0,
-        lora={
-            "enable_lora": True,
-            "lora_rank": 4,
-            "lora_alpha": 8.0,
-            "lora_module_path": ".*decoder.*",
-        },
+        base_emb_dim=32,
+        base_mlp_dim=128,
+        lora=lora_dict,
         steps=1,
+        scan_layers=False,
     )
     (
         _,
@@ -185,15 +196,14 @@ def test_train_step_updates_only_lora_weights(self):
         _,
         train_state,
     ) = setup_train_loop(config, recorder=None)
-    
+
     # Extract initial weights
     initial_model_state = nnx.state(model)
-    
+
     # Run one training step
-    from maxtext.trainers.pre_train import train
     batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
-    
-    new_state, metrics = train.train_step(
+
+    new_state, _ = train.train_step(
         model=nnx.graphdef(train_state),
         config=config,
         state_mesh_shardings=None,
@@ -201,11 +211,11 @@ def test_train_step_updates_only_lora_weights(self):
         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]
@@ -216,7 +226,14 @@ def test_train_step_updates_only_lora_weights(self):
         # 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)
+
+  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)
+
 
 if __name__ == "__main__":
   unittest.main()
-
diff --git a/tests/unit/lora_utils_nnx_test.py b/tests/unit/lora_utils_nnx_test.py
index f3a491d1fd..1da6203b42 100644
--- a/tests/unit/lora_utils_nnx_test.py
+++ b/tests/unit/lora_utils_nnx_test.py
@@ -16,11 +16,15 @@
 Linen regression block at the end."""
 
 import unittest
+from unittest import mock
 
+from flax import nnx
 import jax
 import jax.numpy as jnp
 import numpy as np
 
+from maxtext.utils import lora_utils
+from maxtext.utils import sharding
 from maxtext.utils.lora_utils import (
     apply_lora_on_base_params,
     apply_lora_on_base_params_nnx,
@@ -180,7 +184,11 @@ 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))
+    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(
@@ -217,7 +225,11 @@ def test_numerical_parity_with_linen_apply(self):
     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))
+    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(
@@ -294,31 +306,100 @@ class TestShardingExtractionNnx(unittest.TestCase):
   """Test sharding parameter extraction under different LoRA configurations."""
 
   def test_sharding_extracts_only_lora_params(self):
-    from maxtext.utils import sharding
-    from flax import nnx
     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()
-    graphdef, state = nnx.split(model)
-    
+    _, 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)
+
+
 if __name__ == "__main__":
   unittest.main()
-