[FIX] module_tree tests (#2411)

* FIX moe flag passing not passing nested ci test

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

* fix failed test_model_alignment.py with transformers v5

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

---------

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

Author: ZX-ModelCloud

gptqmodel/models/base.py

@@ -353,7 +353,7 @@ def has_moe_flag(cls, module_spec: str) -> bool:
         return MOE_FLAG.lstrip(":") in flags
 
     @classmethod
-    def _collect_moe_modules_from_tree(cls, tree_node, parent_path="") -> Set[str]:
+    def _collect_moe_modules_from_tree(cls, tree_node, parent_path="", parent_is_moe=False) -> Set[str]:
         """
         Recursively collect all module paths that have the :moe flag.
         Returns a set of full module paths (e.g., "mlp", "mlp.experts", "mlp.shared_experts").
@@ -366,27 +366,26 @@ def _collect_moe_modules_from_tree(cls, tree_node, parent_path="") -> Set[str]:
                 if key == "#":
                     # Recursively process the value if it's a dict
                     if isinstance(value, dict):
-                        moe_modules.update(cls._collect_moe_modules_from_tree(value, parent_path))
+                        moe_modules.update(cls._collect_moe_modules_from_tree(value, parent_path, parent_is_moe))
                     continue
 
                 # Build full path
+                module_name, _ = cls._parse_module_flags(key) if isinstance(key, str) else (key, [])
                 if parent_path:
-                    full_path = f"{parent_path}.{key}"
+                    full_path = f"{parent_path}.{module_name}"
                 else:
-                    full_path = key
+                    full_path = module_name
 
                 # Check if this key has :moe flag
-                if cls.has_moe_flag(key):
-                    # Extract just the module name without flags
-                    module_name, _ = cls._parse_module_flags(key)
-                    if parent_path:
-                        moe_modules.add(f"{parent_path}.{module_name}")
-                    else:
-                        moe_modules.add(module_name)
+                is_moe = cls.has_moe_flag(key) if isinstance(key, str) else False
+                if is_moe or parent_is_moe:
+                    moe_modules.add(full_path)
 
                 # Recursively process nested structures
                 if isinstance(value, (dict, tuple, list)):
-                    moe_modules.update(cls._collect_moe_modules_from_tree(value, full_path.split(":")[0]))
+                    moe_modules.update(
+                        cls._collect_moe_modules_from_tree(value, full_path, parent_is_moe or is_moe)
+                    )
 
         elif isinstance(tree_node, (tuple, list)):
             for item in tree_node:
@@ -397,7 +396,7 @@ def _collect_moe_modules_from_tree(cls, tree_node, parent_path="") -> Set[str]:
                     else:
                         moe_modules.add(module_name)
                 elif isinstance(item, dict):
-                    moe_modules.update(cls._collect_moe_modules_from_tree(item, parent_path))
+                    moe_modules.update(cls._collect_moe_modules_from_tree(item, parent_path, parent_is_moe))
 
         return moe_modules
 

tests/module_tree/test_model_alignment.py

@@ -15,6 +15,7 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "models"))
 from model_test import ModelTest  # noqa: E402
+from torch import nn
 
 
 class TestDots1Struct(ModelTest):
@@ -32,7 +33,11 @@ def test_module_tree_alignment(self):
 
         moe_layer = shell.model.layers[config.first_k_dense_replace]
         self.assertTrue(hasattr(moe_layer.mlp, "experts"))
-        self.assertEqual(len(moe_layer.mlp.experts), config.n_routed_experts)
+        if isinstance(moe_layer.mlp.experts, nn.ModuleList):
+            expert_num = len(moe_layer.mlp.experts)
+        else:
+            expert_num = moe_layer.mlp.experts.num_experts
+        self.assertEqual(expert_num, config.n_routed_experts)
         self.assertTrue(hasattr(moe_layer.mlp, "shared_experts"))
 
         self.assertIn("q_norm:!", Dots1QModel.module_tree[3]["self_attn"])

tests/module_tree/test_moe_flag_parsing.py

@@ -31,20 +31,56 @@ def test_parse_module_flags_combined(self):
         assert "moe" in flags
         assert "!" in flags
 
+    def test_parse_module_flags_with_optional_capture(self):
+        """Test parsing with :? and order variations."""
+        name, flags = BaseQModel._parse_module_flags("mlp:moe:?")
+        assert name == "mlp"
+        assert "moe" in flags
+        assert "?" in flags
+
+        name, flags = BaseQModel._parse_module_flags("mlp:?:moe")
+        assert name == "mlp"
+        assert "moe" in flags
+        assert "?" in flags
+
     def test_has_moe_flag(self):
         """Test MoE flag detection."""
         assert BaseQModel.has_moe_flag("mlp:moe") is True
         assert BaseQModel.has_moe_flag("experts:moe") is True
         assert BaseQModel.has_moe_flag("gate:moe:!") is True
+        assert BaseQModel.has_moe_flag("mlp:moe:?") is True
+        assert BaseQModel.has_moe_flag("mlp:?:moe") is True
         assert BaseQModel.has_moe_flag("gate_proj") is False
         assert BaseQModel.has_moe_flag("gate_proj:!") is False
+        assert BaseQModel.has_moe_flag("mlp:mode:?") is False
+        assert BaseQModel.has_moe_flag("mlp:?:mode") is False
 
     def test_has_moe_flag_non_string(self):
         """Test MoE flag detection with non-string input."""
         assert BaseQModel.has_moe_flag(None) is False
         assert BaseQModel.has_moe_flag(123) is False
         assert BaseQModel.has_moe_flag({}) is False
 
+    def test_collect_moe_modules_with_optional_capture(self):
+        """Test MoE module collection with :? flag in module tree."""
+        module_tree = [
+            "model",
+            "layers",
+            "#",
+            {
+                "mlp:moe:?": {
+                    "experts": {
+                        "#": ("gate_proj:0",),
+                    },
+                },
+            }
+        ]
+
+        MockModel = TestMockMoEModel.create_mock_model_class(module_tree)
+        moe_modules = MockModel.get_moe_modules()
+
+        assert "mlp" in moe_modules
+
 
 class TestMockMoEModel:
     """Mock MoE model for testing."""
@@ -138,6 +174,26 @@ def test_is_moe_module_detection(self):
         assert MockModel.is_moe_module("model.layers.0.self_attn.q_proj") is False
         assert MockModel.is_moe_module("model.layers.0.self_attn") is False
 
+    def test_is_moe_module_with_optional_capture(self):
+        """Test is_moe_module() with :? flag in module tree."""
+        module_tree = [
+            "model",
+            "layers",
+            "#",
+            {
+                "mlp:?:moe": {
+                    "experts": {
+                        "#": ("gate_proj:0",),
+                    },
+                },
+            }
+        ]
+
+        MockModel = TestMockMoEModel.create_mock_model_class(module_tree)
+
+        assert MockModel.is_moe_module("model.layers.0.mlp") is True
+        assert MockModel.is_moe_module("model.layers.0.mlp.experts.3.gate_proj") is True
+
     def test_backward_compatibility_no_moe_flags(self):
         """Test that models without :moe flags still work."""
         module_tree = [