chore: simplify LORA logic and remove debug prints

Author: RexBearIU

src/maxtext/layers/nnx_decoders.py

@@ -589,23 +589,8 @@ def _promote_to_scanned(x):
             curr_params = jax.tree.map(lambda x: _slice_and_unpromote(x, i), params, is_leaf=is_nnx_leaf)
             curr_state = jax.tree.map(lambda x: _slice_and_unpromote(x, i), state, is_leaf=is_nnx_leaf)
 
-            if i == 0:
-                print(f"[DEBUG] Available nodes in Iter 0: {list(curr_params.keys())}")
-
-            print(f"[DEBUG] Iteration {i}: calling layer_fn")
             curr_carry, (out_p, out_o) = layer_fn(curr_carry, (curr_params, curr_state))
 
-            # Inspect results for new LoRA params
-            if i == 0:
-                def find_lora(p, path=""):
-                    if isinstance(p, nnx.State):
-                        for k, v in p.items(): find_lora(v, f"{path}/{k}")
-                    elif isinstance(p, dict):
-                        for k, v in p.items(): find_lora(v, f"{path}/{k}")
-                    elif hasattr(p, "__class__") and p.__class__.__name__ == "LoRAParam":
-                        print(f"[DEBUG] FOUND NEW LORA PARAM at {path}")
-                find_lora(out_p, "params")
-
             # Promote ALL parameters back to rank-3 metadata immediately
             # This ensures they are ready to be stacked correctly.
             out_p = jax.tree.map(_promote_to_scanned, out_p, is_leaf=is_nnx_leaf)
@@ -615,7 +600,6 @@ def find_lora(p, path=""):
             out_other_list.append(out_o)
 
         final_carry = curr_carry
-        print(f"[DEBUG] Loop complete, stacking results and promoting metadata...")
         scanned_params = jax.tree.map(lambda *args: jnp.stack(args), *out_params_list)
         scanned_other = jax.tree.map(lambda *args: jnp.stack(args), *out_other_list)
 
@@ -627,11 +611,12 @@ def _force_promote(x):
             if is_nnx_leaf:
                 metadata = x.get_metadata()
                 updates = {}
+                val_ndim = x.value.ndim
                 for sharding_key in ["sharding", "out_sharding", "sharding_names"]:
                     axes = metadata.get(sharding_key)
                     if isinstance(axes, (jax.sharding.PartitionSpec, tuple, list)):
                         l = list(axes)
-                        if "layers" not in l:
+                        if len(l) < val_ndim and "layers" not in l:
                             pos = min(scan_axis, len(l))
                             l.insert(pos, "layers")
                             updates[sharding_key] = jax.sharding.PartitionSpec(*l) if isinstance(axes, jax.sharding.PartitionSpec) else tuple(l)

src/maxtext/utils/lora_utils.py

@@ -436,31 +436,6 @@ def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperPar
 
   matched_module_paths = []
   sample_module_paths = []
-  found_lora = False
-  seen = set()
-
-  # Truly recursive search to find LoRAParam regardless of NNX registration state
-  def recursive_find_lora(obj):
-      nonlocal found_lora
-      if found_lora or id(obj) in seen: return
-      seen.add(id(obj))
-      
-      if hasattr(obj, "__class__") and obj.__class__.__name__ == "LoRAParam":
-          found_lora = True
-          return
-      
-      if hasattr(obj, "__dict__"):
-          for k, v in obj.__dict__.items():
-              if not k.startswith("__"):
-                  recursive_find_lora(v)
-      elif isinstance(obj, (dict, list, tuple)):
-          items = obj.values() if isinstance(obj, dict) else obj
-          for v in items: recursive_find_lora(v)
-
-  recursive_find_lora(lora_model)
-
-  if found_lora:
-    return
 
   for path, _ in nnx.iter_graph(lora_model):
     module_path = "/".join(str(p) for p in path)
@@ -578,6 +553,12 @@ def patched_get_or_create_lora_params(*, name, rule, a_shape, b_shape, a_shardin
         b_sharding_transpose=b_sharding_transpose,
     )
 
+    # Ensure they are specifically LoRAParam, not just generic Param or Variable
+    if hasattr(lora_a, "value") and hasattr(lora_a, "get_metadata"):
+        lora_a = nnx.LoRAParam(lora_a.value, **lora_a.get_metadata())
+    if hasattr(lora_b, "value") and hasattr(lora_b, "get_metadata"):
+        lora_b = nnx.LoRAParam(lora_b.value, **lora_b.get_metadata())
+
     # Force registration on the current module
     module = flax_util.get_current_module()
     if isinstance(module, nnx.Module):
@@ -624,20 +605,12 @@ def apply_lora_to_model(
   finally:
     model.decoder.disable_quant_stats_update = False
 
-  # Important: use the NEW model returned by Qwix!
+  # Important: Qwix dynamically swaps the __class__ of the model, which breaks nnx.iter_graph
+  # We must restore the original unquantized class type for Tunix to recognize the module correctly.
+  if hasattr(lora_model, "_unquantized_type"):
+      lora_model.__class__ = getattr(lora_model, "_unquantized_type")
+
   model = lora_model
-  
-  # Check if we can find lora in this model immediately
-  temp_found = []
-  def quick_check(obj, path=""):
-      if len(temp_found) > 0: return
-      if hasattr(obj, "__class__") and obj.__class__.__name__ == "LoRAParam":
-          temp_found.append(path)
-      if hasattr(obj, "__dict__"):
-          for k,v in obj.__dict__.items():
-              if not k.startswith("__"): quick_check(v, f"{path}/{k}")
-  quick_check(model, "root")
-  print(f"[DEBUG] Quick check for LoRA in lora_model: {temp_found}")
 
   def rank_consistent_spec(spec, shape):
     if spec is None: return None
@@ -654,7 +627,7 @@ def rank_consistent_spec(spec, shape):
 
   if mesh is not None:
     with mesh, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
-      graph_def, state = nnx.split(lora_model)
+      graph_def, state = nnx.split(model)
 
       def fix_metadata(x):
         if hasattr(x, "get_metadata") and hasattr(x, "replace"):