Fix QLoRA NNX issues

.github/workflows/ci_pipeline.yml

@@ -225,8 +225,7 @@ jobs:
     needs: [gate_test_run]
     if: |
       always() &&
-      needs.gate_test_run.result == 'success' &&
-      github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
+      needs.gate_test_run.result == 'success'
     uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false

src/maxtext/inference/maxengine/maxengine.py

@@ -485,7 +485,11 @@ def _overlay(dst, src):
           lambda x: jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(None, *x.spec)),
           self.prefill_kv_cache_shardings,
       )
-      self.prefill_kv_cache_shardings = {"decoder": {"layers": self.prefill_kv_cache_shardings["decoder"]["layers"][0]}}
+      if "layers" in self.prefill_kv_cache_shardings["decoder"]:
+        first_layer_sharding = self.prefill_kv_cache_shardings["decoder"]["layers"][0]
+      else:
+        first_layer_sharding = self.prefill_kv_cache_shardings["decoder"]["layers_0"]
+      self.prefill_kv_cache_shardings = {"decoder": {"layers": first_layer_sharding}}
     # scan_layers=True is already stacked on axis 0; shardings stay as-is and stack/unstack are no-ops.
     # AR-mode abstract model so axis names use CACHE_BATCH (not CACHE_BATCH_PREFILL);
     # bulk_insert / _insert_jit search for "cache_batch" in the per-leaf logical axes.
@@ -609,10 +613,16 @@ def _maybe_stack_prefill_result_cache(self, cache):
       if self.config.scan_layers:
         # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack.
         return cache
-      # scan_layers=False: stack the per-layer subtrees under decoder/layers into one
+      # scan_layers=False: stack the per-layer subtrees under decoder into one
       # subtree with a leading layer axis (matching the scan_layers=True shape).
-      layers = cache["decoder"]["layers"]
-      stacked = jax.tree.map(lambda *c: jnp.stack(c), *[layers[i] for i in range(self.config.num_decoder_layers)])
+      if "layers" in cache["decoder"]:
+        layers = cache["decoder"]["layers"]
+        layer_cache = [layers[i] for i in range(self.config.num_decoder_layers)]
+      else:
+        layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
+        layer_cache = [cache["decoder"][key] for key in layer_keys]
+
+      stacked = jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
       return {"decoder": {"layers": stacked}}
 
     layer_keys = []
@@ -635,8 +645,10 @@ def _maybe_unstack_prefill_result_cache(self, cache):
         return cache
       # scan_layers=False: split the leading layer axis back into per-layer subtrees.
       stacked = cache["decoder"]["layers"]
-      layers = {i: jax.tree.map(lambda x, i=i: x[i], stacked) for i in range(self.config.num_decoder_layers)}
-      return {"decoder": {"layers": layers}}
+      res_cache = {"decoder": {}}
+      for i in range(self.config.num_decoder_layers):
+        res_cache["decoder"][f"layers_{i}"] = jax.tree.map(lambda x, i=i: x[i], stacked)
+      return res_cache
 
     flat_cache, treedef = jax.tree.flatten(cache)
     layer_cache = [jax.tree.unflatten(treedef, flat_cache_vars) for flat_cache_vars in zip(*flat_cache, strict=True)]

src/maxtext/layers/nnx_decoders.py

@@ -498,13 +498,13 @@ def _init_pipeline_deepseek(self, decoder_block_classes, rngs):
     else:
       self.num_dense_layers = config.first_num_dense_layers
       for i in range(self.num_dense_layers):
-        self._create_and_register_named_layer(dense_cls, rngs, "dense_layers", i)
+        self._create_and_register_layer(dense_cls, rngs, "dense_layers", i)
       self.num_moe_outside_pipeline = (
           config.num_decoder_layers - config.first_num_dense_layers
       ) - config.pipeline_parallel_layers
       if self.num_moe_outside_pipeline > 0:
         for i in range(self.num_moe_outside_pipeline):
-          self._create_and_register_named_layer(moe_cls, rngs, "moe_layers_outside_pipeline", i)
+          self._create_and_register_layer(moe_cls, rngs, "moe_layers_outside_pipeline", i)
 
   def _init_pipeline_generic(self, decoder_block_classes, rngs):
     """Initializes generic decoder layers outside pipeline."""
@@ -522,7 +522,7 @@ def _init_pipeline_generic(self, decoder_block_classes, rngs):
       else:
         self.num_layers_outside_pipeline = remaining_layers
         for i in range(self.num_layers_outside_pipeline):
-          self._create_and_register_named_layer(base_cls, rngs, "layers_outside_pipeline", i)
+          self._create_and_register_layer(base_cls, rngs, "layers_outside_pipeline", i)
 
   def _init_scanned_layers(self, decoder_block_classes, rngs, mesh):
     """Initializes decoder layers with scanning (non-pipeline)."""
@@ -689,12 +689,9 @@ def _init_scanned_generic(self, decoder_block_classes, rngs):
           rngs=rngs,
           **layer_kwargs,
       )
-    else:
-      self.layers = nnx.List([])
 
   def _init_sequential_layers(self, decoder_block_classes, rngs):
     """Initializes decoder layers sequentially (no scanning)."""
-    self.layers = nnx.List([])
 
     if self.is_deepseek:
       self._init_sequential_deepseek(decoder_block_classes, rngs)
@@ -706,9 +703,9 @@ def _init_sequential_deepseek(self, decoder_block_classes, rngs):
     config = self.config
     dense_cls, moe_cls = decoder_block_classes
     for i in range(config.first_num_dense_layers):
-      self._create_and_register_layer(dense_cls, rngs, "dense_layer", i)
+      self._create_and_register_layer(dense_cls, rngs, "dense_layers", i)
     for i in range(config.num_decoder_layers - config.first_num_dense_layers):
-      self._create_and_register_layer(moe_cls, rngs, "moe_layer", i)
+      self._create_and_register_layer(moe_cls, rngs, "moe_layers", i)
 
   def _init_sequential_generic(self, decoder_block_classes, rngs):
     """Initializes sequential generic decoder layers with per-architecture layer_kwargs."""
@@ -749,7 +746,6 @@ def _init_gemma4_small_layers(self, rngs):
     ``_create_and_register_layer``.
     """
     cfg = self.config
-    self.layers = nnx.List([])
     # Only register the PLE submodule when it exists (mirrors the optional position_embedder
     # pattern); assigning None first would make nnx treat the attribute as static.
     if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0:
@@ -767,7 +763,6 @@ def _init_gemma4_small_layers(self, rngs):
           rngs=rngs,
       )
       setattr(self, f"layers_{lyr}", layer)
-      self.layers.append(layer)
 
   def _get_pipeline_stage_module(self, decoder_blocks, rngs):
     """Retrieves the wrapper module formatted for single pipeline stage execution."""
@@ -797,14 +792,7 @@ def _get_pipeline_stage_module(self, decoder_blocks, rngs):
     )
 
   def _create_and_register_layer(self, layer_cls, rngs, base_name, i, **layer_kwargs):
-    attr_name = f"{base_name}_{i}"
-    layer = self._create_single_layer(layer_cls, rngs, **layer_kwargs)
-    setattr(self, attr_name, layer)
-    self.layers.append(layer)
-
-  def _create_and_register_named_layer(self, layer_cls, rngs, base_name, i, **layer_kwargs):
-    """Creates a layer registered ONLY via named attribute. Used by pipeline-outside paths
-    to avoid double-registration when self.layers list is also tracked elsewhere."""
+    """Creates a layer registered ONLY via named attribute."""
     attr_name = f"{base_name}_{i}"
     layer = self._create_single_layer(layer_cls, rngs, **layer_kwargs)
     setattr(self, attr_name, layer)
@@ -1846,22 +1834,19 @@ def __call__(
       else:
         prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg)
 
-        # Hoisted function to preserve XLA cache ID
-        def pure_layer_fn(graphdef, state_in, y_in, kv_in):
-
-          if cfg.parameter_memory_host_offload:
-            state_in = jax.tree.map(
-                lambda x: jax.device_put(x, max_utils.device_space()),
-                state_in,
-            )
-
-          merged_layer = nnx.merge(graphdef, state_in)
-          out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs)
-          return out_y, out_kv, nnx.state(merged_layer)
-
-        checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
+        # Define dynamic_graph_init and updated_graphdefs
+        dynamic_graph_init = bool(getattr(self, "disable_quant_stats_update", False))
+        updated_graphdefs = [None] * cfg.num_decoder_layers
 
-        for lyr, layer in enumerate(self.layers):
+        for lyr in range(cfg.num_decoder_layers):
+          if self.is_deepseek:
+            if lyr < cfg.first_num_dense_layers:
+              layer = getattr(self, f"dense_layers_{lyr}")
+            else:
+              moe_idx = lyr - cfg.first_num_dense_layers
+              layer = getattr(self, f"moe_layers_{moe_idx}")
+          else:
+            layer = getattr(self, f"layers_{lyr}", self.layers[lyr] if hasattr(self, "layers") and self.layers else None)
           graphdef, state = nnx.split(layer)
           if kv_caches is not None:
             if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5):
@@ -1881,8 +1866,46 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in):
           if input_tokens is not None:
             layer_kwargs["decoder_input_tokens"] = input_tokens
 
-          y, kv_cache, new_state = checkpointed_fn(graphdef, state, y, kv_cache)
-          nnx.update(layer, new_state)
+          # Define pure_layer_fn locally to capture 'lyr' and 'updated_graphdefs'
+          def pure_layer_fn_local(graphdef_in, state_in, y_in, kv_in, lyr=lyr):
+            if cfg.parameter_memory_host_offload:
+              state_in = jax.tree.map(
+                  lambda x: jax.device_put(x, max_utils.device_space()),
+                  state_in,
+              )
+            merged_layer = nnx.merge(graphdef_in, state_in)
+            out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs)
+            state_out = nnx.state(merged_layer)
+
+            if dynamic_graph_init:
+              new_graphdef, _, _ = nnx.split(merged_layer, nnx.Param, ...)
+              updated_graphdefs[lyr] = new_graphdef
+
+            return out_y, out_kv, state_out
+
+          checkpointed_fn_local = jax.checkpoint(pure_layer_fn_local, policy=policy, prevent_cse=prevent_cse)
+
+          if cfg.remat_policy != "none":
+            y, kv_cache, new_state = checkpointed_fn_local(graphdef, state, y, kv_cache)
+          else:
+            y, kv_cache, new_state = pure_layer_fn_local(graphdef, state, y, kv_cache)
+
+          if dynamic_graph_init:
+            new_params, new_rest = new_state.split(nnx.Param, ...)
+            new_layer = nnx.merge(updated_graphdefs[lyr], new_params, new_rest)
+            if self.is_deepseek:
+              if lyr < cfg.first_num_dense_layers:
+                setattr(self, f"dense_layers_{lyr}", new_layer)
+              else:
+                moe_idx = lyr - cfg.first_num_dense_layers
+                setattr(self, f"moe_layers_{moe_idx}", new_layer)
+            else:
+              if hasattr(self, "layers") and self.layers:
+                self.layers[lyr] = new_layer
+              else:
+                setattr(self, f"layers_{lyr}", new_layer)
+          else:
+            nnx.update(layer, new_state)
 
           if kv_caches is not None and kv_cache is not None:
             if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5):
@@ -2113,7 +2136,7 @@ def _apply_gemma4_small_layers(
     cache_index_of = gemma4_small.kv_cache_slot_map(layer_types, num_kv_shared)
 
     for lyr in range(cfg.num_decoder_layers):
-      layer = self.layers[lyr]
+      layer = getattr(self, f"layers_{lyr}")
       donor_idx = gemma4_small.kv_donor_layer_idx(lyr, layer_types, num_kv_shared)
       is_donor = gemma4_small.is_kv_donor_layer(lyr, layer_types, num_kv_shared)
 
@@ -2158,6 +2181,14 @@ def _apply_gemma4_small_layers(
 
     return y, kv_caches
 
+  def __getattr__(self, name):
+    if name == "layers" and hasattr(self, "layers_0"):
+      return [getattr(self, f"layers_{i}") for i in range(self.config.num_decoder_layers) if hasattr(self, f"layers_{i}")]
+    if hasattr(super(), "__getattr__"):
+      return super().__getattr__(name)
+    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+
 
 def decoder_as_linen(
     config: Config,

src/maxtext/utils/lora_utils.py

@@ -433,10 +433,10 @@ def _get_lora_module_path(mt_config: pyconfig.HyperParameters) -> str:
 
   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/...').
-  optional_layer_index = "(?:[0-9]+/)?"
-  final_path = str(raw_path).replace("layers/", f"layers/{optional_layer_index}")
+  # This regex makes the layer index optional, matching scanned, unscanned named (layers_0),
+  # and unscanned index (layers/0) layer paths.
+  layer_pattern = r"layers(?:_[0-9]+|/[0-9]+)?/"
+  final_path = str(raw_path).replace("layers/", layer_pattern)
 
   max_logging.log(f"Using lora_module_path: {final_path}")
   return final_path
@@ -579,6 +579,7 @@ def apply_lora_to_model(
     mt_config: pyconfig.HyperParameters,
 ) -> nnx.Module:
   """Optionally applies LoRA/QLoRA to a MaxText model using Qwix."""
+  # pylint: disable=protected-access
   # Skip Qwix LoRA if MaxText LoRA adapters are loaded
   if mt_config.lora_input_adapters_path:
     max_logging.log("MaxText LoRA adapters loaded, skipping Qwix LoRA application")
@@ -587,6 +588,94 @@ def apply_lora_to_model(
   if not mt_config.lora.enable_lora:
     return model
 
+  # Monkeypatch qwix.LoraProvider.dot_general to fix find_param with PTQ
+  if not hasattr(qwix.LoraProvider, "_patched_for_maxtext"):
+    # pylint: disable=import-outside-toplevel
+    from qwix._src.providers.lora import LoraProvider, LoraRule, _create_lora_layer_shapes, _get_or_create_lora_params, _compute_lora_delta
+    from qwix._src.providers import ptq
+
+    try:
+      from qwix._src.utils import flax_util
+    except ImportError:
+      from qwix._src import flax_util
+    # pylint: enable=import-outside-toplevel
+
+    def patched_dot_general(
+        self,
+        lhs: jax.Array,
+        rhs: jax.Array,
+        dimension_numbers: jax.lax.DotDimensionNumbers,
+        precision: jax.lax.PrecisionLike = None,
+        preferred_element_type: jax.typing.DTypeLike | None = None,
+        out_sharding: jax.sharding.NamedSharding | None = None,
+    ) -> jax.Array:
+      weight_name = flax_util.find_param(rhs)
+
+      res = ptq.PtqProvider.dot_general(
+          self,
+          lhs,
+          rhs,
+          dimension_numbers,
+          precision,
+          preferred_element_type,
+          out_sharding=out_sharding,
+      )
+
+      rule, _ = self._get_current_rule_and_op_id("dot_general", repeated_call=True)
+      if not isinstance(rule, LoraRule):
+        return res
+
+      if weight_name is None:  # rhs is not a weight.
+        return res
+
+      (lhs_ca, rhs_ca), (lhs_ba, rhs_ba) = dimension_numbers
+
+      rhs_ra = (*(i for i in range(rhs.ndim) if i not in rhs_ca and i not in rhs_ba),)
+
+      contract_shape = (*(rhs.shape[i] for i in rhs_ca),)
+      batch_shape = (*(rhs.shape[i] for i in rhs_ba),)
+      remain_shape = (*(rhs.shape[i] for i in rhs_ra),)
+
+      a_shape, a_sharding_transpose, b_shape, b_sharding_transpose = _create_lora_layer_shapes(
+          rhs_ca,
+          rhs_ba,
+          rhs_ra,
+          contract_shape,
+          batch_shape,
+          remain_shape,
+          rule.rank,
+      )
+
+      lora_a, lora_b = _get_or_create_lora_params(
+          name=weight_name,
+          rule=rule,
+          a_shape=a_shape,
+          b_shape=b_shape,
+          a_sharding_transpose=a_sharding_transpose,
+          b_sharding_transpose=b_sharding_transpose,
+      )
+
+      if rule.dropout > 0:
+        lhs = nnx.Dropout(rule.dropout, deterministic=False)(lhs, rngs=flax_util.make_rng("dropout"))
+
+      delta = _compute_lora_delta(
+          lhs,
+          lora_a,
+          lora_b,
+          lhs_ca,
+          lhs_ba,
+          contract_shape,
+          batch_shape,
+          remain_shape,
+          rule.rank,
+          precision=precision,
+      )
+
+      return res + delta * (rule.alpha / rule.rank)
+
+    LoraProvider.dot_general = patched_dot_general
+    LoraProvider._patched_for_maxtext = True
+
   # Dynamically detect and set LoRA rank before model creation if restoring
 
   lora_provider = _build_lora_provider(mt_config)