feat: Gemma4 LoRA Extension

Author: RexBearIU

src/maxtext/configs/post_train/lora_module_path.yml

@@ -21,6 +21,7 @@ mistral: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
 deepseek2: "decoder/(dense_layers|moe_stack)/self_attention/(query|out|wkv_a|wkv_b)|decoder/(dense_layers|moe_stack)/(mlp|shared_experts)/(wi_0|wi_1|wo)"
 gemma2: "decoder/layers/(self_attention_local|self_attention_global)/(query|key|value|out)|decoder/layers/(mlp_local|mlp_global)/(wi_0|wi_1|wo)"
 gemma3: "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo|gate|up|down))"
+gemma4: "decoder/(scanned_blocks|layers_remainder)/layers.*/.*(self_attention/(query|key|value|out)|mlp/.*(wi_0|wi_1|wo|shared_experts/(wi_0|wi_1|wo)))"
 olmo3: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
 gpt3: "decoder/layers/(self_attention/(qkv_proj|out)|mlp/(wi|wo))"
 

src/maxtext/trainers/post_train/sft/train_sft.py

@@ -264,9 +264,14 @@ def setup_trainer_state(mt_config, goodput_recorder=None):
 def train_model(mt_config, trainer, mesh):
   """Runs the SFT training loop in Tunix."""
   with mesh, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+    # Disable NNX graph caching for MoE models (where experts > 1) to allow
+    # necessary dynamic metadata synchronization during forward passes (e.g., in jax.lax.scan).
+    enable_nnx_cache = mt_config.num_experts <= 1
+
     trainer.train(
         trainer.data_hooks.train_data_iterator,
         trainer.data_hooks.eval_data_iterator,
+        cache_nnx_graph=enable_nnx_cache,
     )
   return trainer
 

src/maxtext/utils/lora_utils.py

@@ -27,7 +27,6 @@
 import jax
 import jax.numpy as jnp
 from orbax import checkpoint as ocp
-import qwix
 
 from maxtext.common import checkpointing
 from maxtext.configs import pyconfig
@@ -408,8 +407,10 @@ def _get_lora_module_path(mt_config: pyconfig.HyperParameters) -> str:
   return final_path
 
 
-def _build_lora_provider(mt_config: pyconfig.HyperParameters) -> qwix.LoraProvider:
+def _build_lora_provider(mt_config: pyconfig.HyperParameters) -> "qwix.LoraProvider":
   """Builds a Qwix LoRA provider from MaxText LoRA settings."""
+  import qwix  # pylint: disable=import-outside-toplevel
+
   lora_module_path = _get_lora_module_path(mt_config)
   lora_kwargs = {
       "module_path": lora_module_path,
@@ -495,6 +496,8 @@ def apply_lora_to_model(
   model_rngs = getattr(model.decoder, "rngs", None)
   decoder_input_tokens, decoder_positions = _prepare_dummy_inputs()
 
+  import qwix  # pylint: disable=import-outside-toplevel
+
   lora_model = qwix.apply_lora_to_model(
       model,
       lora_provider,