[FIX] qwen3_5_moe / llama4 / qwen2_moe / qwen3_next awq layer grouping (#2634)

* fix qwen3.5 moe awq layer grouping

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

* fix llama4/qwen2_moe/qwen3_next awq layer grouping

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

---------

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

Author: ZX-ModelCloud

gptqmodel/models/definitions/llama4.py

@@ -32,10 +32,10 @@ class Llama4QModel(BaseQModel):
             "self_attn": ("q_proj:0", "k_proj:0", "v_proj:0", "o_proj:1"),
             "post_attention_layernorm": ("post_attention_layernorm:!",),
             "feed_forward:moe": {
-                "experts": {
+                "experts:0": {
                     "#": ("gate_proj:0", "up_proj:0", "down_proj:1"),
                 },
-                "shared_expert": ("gate_proj:0", "up_proj:0", "down_proj:1"),
+                "shared_expert:0": ("gate_proj:0", "up_proj:0", "down_proj:1"),
             },
         }
     ]

gptqmodel/models/definitions/qwen2_moe.py

@@ -35,17 +35,3 @@ class Qwen2MoeQModel(BaseQModel):
             },
         }
     ]
-
-    # module_tree_overrides = {
-    #     METHOD.AWQ: [
-    #         {
-    #             "mlp:moe:?": {
-    #                 "gate": ("gate:!",),
-    #                 "shared_expert": None,
-    #                 "experts": {
-    #                     "#": ("gate_proj:0", "up_proj:0", "down_proj:1"),
-    #                 },
-    #             },
-    #         }
-    #     ]
-    # }

gptqmodel/models/definitions/qwen3_5_moe.py

@@ -55,10 +55,10 @@ class Qwen3_5_MoeQModel(BaseQModel):
             "mlp:moe:?": {
                 "gate": ("gate:!",),  # <-- 0.5MB per layer. Not worth quantizing
                 "shared_expert_gate": ("shared_expert_gate:!",),
-                "experts": {
+                "experts:0": {
                     "#": ("gate_proj:0", "up_proj:0", "down_proj:1"),
                 },
-                "shared_expert": ("gate_proj:0", "up_proj:0", "down_proj:1"),
+                "shared_expert:0": ("gate_proj:0", "up_proj:0", "down_proj:1"),
             },
         }
     ]

gptqmodel/models/definitions/qwen3_next.py

@@ -44,10 +44,10 @@ class Qwen3NextGPTQ(BaseQModel):
                 # MoE router + shared expert (Qwen3NextSparseMoeBlock)
                 "gate": ("gate:!",),  # router gate linear
                 "shared_expert_gate": ("shared_expert_gate:!",), # <-- single (1, N) logic projections should not be quantized
-                "shared_expert": ("gate_proj:0", "up_proj:0", "down_proj:1"),
+                "shared_expert:0": ("gate_proj:0", "up_proj:0", "down_proj:1"),
 
                 # Experts list with dynamic index
-                "experts": {
+                "experts:0": {
                     "#": ("gate_proj:0", "up_proj:0", "down_proj:1"),
                 },
             },

tests/module_tree/test_subset.py

@@ -27,6 +27,7 @@
 from gptqmodel.looper.named_module import NamedModule
 from gptqmodel.looper.stage_subset import build_subset_plan, run_subset_stage
 from gptqmodel.models.definitions.qwen2_moe import Qwen2MoeQModel
+from gptqmodel.models.definitions.qwen3_5_moe import Qwen3_5_MoeQModel
 from gptqmodel.models.definitions.qwen3_moe import Qwen3MoeQModel
 from gptqmodel.nn_modules.hooked_linear import replace_module_with_hooked_legacy
 from gptqmodel.quantization import FORMAT, METHOD
@@ -111,6 +112,20 @@ def test_qwen2_moe_shared_expert_merges_with_experts():
     assert len(expert_gate_blocks) == 1
 
 
+def test_qwen3_5_moe_shared_expert_merges_with_experts():
+    blocks = Qwen3_5_MoeQModel.build_layer_modules(Qwen3_5_MoeQModel.module_tree)
+    print("blocks",blocks)
+    gate_block = next(block for block in blocks if "mlp.shared_expert.gate_proj" in block)
+    assert "mlp.experts.{expert_index}.gate_proj" in gate_block
+    assert "mlp.experts.{expert_index}.up_proj" in gate_block
+
+    down_block = next(block for block in blocks if "mlp.shared_expert.down_proj" in block)
+    assert "mlp.experts.{expert_index}.down_proj" in down_block
+
+    expert_gate_blocks = [block for block in blocks if "mlp.experts.{expert_index}.gate_proj" in block]
+    assert len(expert_gate_blocks) == 1
+
+
 def test_awq_processor_enables_subset_early_stop():
     calibration = [{"input_ids": torch.tensor([1, 2, 3])}]
     qcfg = _make_quant_config()