Add unit tests for Gemma 4 LoRA targeting and NNX graph caching

Author: RexBearIU

src/maxtext/configs/post_train/lora_module_path.yml

@@ -21,7 +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/(scanned_blocks|layers_remainder|layers)/(self_attention_local|self_attention_global)/(query|key|value|out)|decoder/(scanned_blocks|layers_remainder|layers)/(mlp_local|mlp_global)/(wi_0|wi_1|wo)"
 gemma3: "decoder/(scanned_blocks|layers_remainder|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)))"
+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))"
 

tests/post_training/unit/lora_utils_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Tests for Qwix LoRA utils in lora_utils.py"""
+import re
 import sys
 import unittest
 from unittest import mock
@@ -84,6 +85,14 @@ def test_get_lora_module_path(self):
         "decoder/layers/(?:[0-9]+/)?.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))",
     )
 
+    mock_config.model_name = "gemma4-9b"
+    path = lora_utils._get_lora_module_path(mock_config)
+    self.assertEqual(
+        path,
+        "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)))",
+    )
+
     mock_config.model_name = "unknown_model"
     # Ensure lora.lora_module_path is still empty string to trigger fallback
     mock_config.lora.lora_module_path = ""
@@ -262,6 +271,66 @@ def test_restore_lora_from_path(self):
           self.assertTrue(kwargs["args"].partial_restore)
         mock_update.assert_called_once()
 
+  def test_gemma4_lora_path_matching(self):
+    """Test that the Gemma4 LoRA regex correctly matches all expected parameter paths."""
+    mock_config = mock.MagicMock(spec=pyconfig.HyperParameters)
+    mock_config.lora = mock.MagicMock()
+    mock_config.lora.lora_module_path = ""
+    mock_config.model_name = "gemma4-9b"
+
+    path_regex = lora_utils._get_lora_module_path(mock_config)
+    compiled = re.compile(path_regex)
+
+    # Expected matching paths:
+    matching_paths = [
+        # Scan layers = True, Dense/MoE attention
+        "decoder/scanned_blocks/layers/self_attention/query/kernel",
+        "decoder/scanned_blocks/layers/self_attention/key/kernel",
+        "decoder/scanned_blocks/layers/self_attention/value/kernel",
+        "decoder/scanned_blocks/layers/self_attention/out/kernel",
+        # Scan layers = True, Dense MLP
+        "decoder/scanned_blocks/layers/mlp/wi_0/kernel",
+        "decoder/scanned_blocks/layers/mlp/wi_1/kernel",
+        "decoder/scanned_blocks/layers/mlp/wo/kernel",
+        # Scan layers = True, MoE MLP
+        "decoder/scanned_blocks/layers/mlp/shared_experts/wi_0/kernel",
+        "decoder/scanned_blocks/layers/mlp/shared_experts/wi_1/kernel",
+        "decoder/scanned_blocks/layers/mlp/shared_experts/wo/kernel",
+        # Scan layers = False, Dense/MoE attention
+        "decoder/layers_remainder/layers/0/self_attention/query/kernel",
+        "decoder/layers_remainder/layers/0/self_attention/key/kernel",
+        "decoder/layers_remainder/layers/0/self_attention/value/kernel",
+        "decoder/layers_remainder/layers/0/self_attention/out/kernel",
+        # Scan layers = False, Dense MLP
+        "decoder/layers_remainder/layers/0/mlp/wi_0/kernel",
+        "decoder/layers_remainder/layers/0/mlp/wi_1/kernel",
+        "decoder/layers_remainder/layers/0/mlp/wo/kernel",
+        # Scan layers = False, MoE MLP
+        "decoder/layers_remainder/layers/0/mlp/shared_experts/wi_0/kernel",
+        "decoder/layers_remainder/layers/0/mlp/shared_experts/wi_1/kernel",
+        "decoder/layers_remainder/layers/0/mlp/shared_experts/wo/kernel",
+        # No scanned_blocks/layers_remainder prefix (e.g. fallback or direct structure)
+        "decoder/layers/0/self_attention/query/kernel",
+        "decoder/layers/0/mlp/wi_0/kernel",
+        "decoder/layers/layers/0/mlp/shared_experts/wi_0/kernel",
+    ]
+
+    for path in matching_paths:
+      self.assertTrue(compiled.search(path), f"Failed to match valid path: {path}")
+
+    # Expected non-matching paths (e.g. layernorm, embedding):
+    non_matching_paths = [
+        "decoder/scanned_blocks/layers/pre_self_attention_norm/scale",
+        "decoder/scanned_blocks/layers/post_self_attention_norm/scale",
+        "decoder/layers_remainder/layers/0/pre_self_attention_norm/scale",
+        "decoder/layers_remainder/layers/0/post_self_attention_norm/scale",
+        "decoder/final_norm/scale",
+        "token_embedder/embedding",
+    ]
+
+    for path in non_matching_paths:
+      self.assertFalse(compiled.search(path), f"Incorrectly matched invalid path: {path}")
+
 
 if __name__ == "__main__":
   unittest.main()

tests/post_training/unit/train_sft_test.py

@@ -15,6 +15,7 @@
 """Unit tests for train_sft.py."""
 
 import unittest
+from unittest import mock
 from types import SimpleNamespace
 import pytest
 
@@ -42,6 +43,48 @@ def test_validate_config_invalid_offload(self):
     with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"):
       train_sft.validate_config(config)
 
+  @pytest.mark.cpu_only
+  def test_train_model_caching_moe(self):
+    """Test that NNX graph caching is disabled for MoE models (num_experts > 1)."""
+    mt_config = SimpleNamespace(
+        logical_axis_rules=[],
+        num_experts=8,
+    )
+    trainer = mock.MagicMock()
+    trainer.data_hooks.train_data_iterator = "train_iter"
+    trainer.data_hooks.eval_data_iterator = "eval_iter"
+    mesh = mock.MagicMock()
+
+    with mock.patch("jax.set_mesh"):
+      train_sft.train_model(mt_config, trainer, mesh)
+
+    trainer.train.assert_called_once_with(
+        "train_iter",
+        "eval_iter",
+        cache_nnx_graph=False,
+    )
+
+  @pytest.mark.cpu_only
+  def test_train_model_caching_dense(self):
+    """Test that NNX graph caching is enabled for dense models (num_experts <= 1)."""
+    mt_config = SimpleNamespace(
+        logical_axis_rules=[],
+        num_experts=1,
+    )
+    trainer = mock.MagicMock()
+    trainer.data_hooks.train_data_iterator = "train_iter"
+    trainer.data_hooks.eval_data_iterator = "eval_iter"
+    mesh = mock.MagicMock()
+
+    with mock.patch("jax.set_mesh"):
+      train_sft.train_model(mt_config, trainer, mesh)
+
+    trainer.train.assert_called_once_with(
+        "train_iter",
+        "eval_iter",
+        cache_nnx_graph=True,
+    )
+
 
 if __name__ == "__main__":
   unittest.main()