diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index 5e2237aa87..a5a803acc4 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.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 diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 095d3ea48a..80ff14b9a2 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/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)] diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 528523e438..30f8eec6f7 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/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, diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 1aa443f093..2c6ee633fc 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/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) diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index e910995542..86d7a30160 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/tests/post_training/unit/lora_utils_test.py @@ -49,10 +49,10 @@ "base_num_decoder_layers": 1, "attention": "dot_product", "max_target_length": 8, - "base_emb_dim": 128, + "base_emb_dim": 256, "base_num_query_heads": 2, "base_num_kv_heads": 2, - "base_mlp_dim": 256, + "base_mlp_dim": 512, "max_prefill_predict_length": 4, "model_name": "llama2-7b", "enable_nnx": True, @@ -88,7 +88,7 @@ def test_get_lora_module_path(self): path = lora_utils._get_lora_module_path(mock_config) self.assertEqual( path, - "decoder/layers/(?:[0-9]+/)?.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))", + "decoder/layers(?:_[0-9]+|/[0-9]+)?/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))", ) mock_config.model_name = "gemma4-9b" @@ -106,7 +106,7 @@ def test_get_lora_module_path(self): # Fallback to default self.assertEqual( path, - "decoder/layers/(?:[0-9]+/)?.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))", + "decoder/layers(?:_[0-9]+|/[0-9]+)?/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))", ) mock_config.lora.lora_module_path = "custom/path" @@ -244,6 +244,15 @@ def _run_apply_lora_test( # Verify it IS now LoRA enabled self.assertTrue(lora_utils.is_lora_enabled(lora_model)) + # Verify quantization if weight_qtype is set + flat_state_paths = [".".join(str(p) for p in path) for path, _ in state.flat_state()] + if weight_qtype is not None: + self.assertTrue(any("qvalue" in path for path in flat_state_paths), "Expected quantized weights (qvalue) in state") + else: + self.assertFalse( + any("qvalue" in path for path in flat_state_paths), "Did not expect quantized weights (qvalue) in state" + ) + # Test fit for PeftTrainer trainer_cfg = peft_trainer.TrainingConfig(eval_every_n_steps=10) optimizer = optax.adam(1e-4) diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 4585c9c3ee..f90d9ac34a 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -799,6 +799,7 @@ def test_gemma4_small_decoder(self): "hidden_size_per_layer_input=128", "vocab_size_per_layer_input=256", "vocab_size=256", + "max_target_length=128", ], override_model_config=True, ) @@ -872,6 +873,7 @@ def test_gemma4_small_decoder_with_mock_cache_and_ple(self): "hidden_size_per_layer_input=128", "vocab_size_per_layer_input=256", "vocab_size=256", + "max_target_length=128", ], override_model_config=True, )