[GraniteMoE] Load split per-expert quantized MoE weights (#37679)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

test/registered/unit/models/test_granitemoe_split_expert_loading.py

@@ -0,0 +1,195 @@
+"""Unit tests for granitemoe_load_split_experts.
+
+Compressed-tensors checkpoints (llmcompressor) store MoE experts one tensor per
+expert per projection -- `experts.<id>.{gate,up,down}_proj.{weight,weight_scale}`
+-- while the unquantized HF checkpoint packs them into `input_linear` /
+`output_linear`. Only the packed layout was recognised, so every split expert
+tensor fell through to a `logger.warning(...not found in params_dict)` and was
+silently dropped: the server started normally and then emitted garbage
+("capital capital capital..." instead of " Paris.").
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=6, suite="base-a-test-cpu")
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+from sglang.srt.models.granitemoe import (
+    _is_packed_expert,
+    granitemoe_load_split_experts,
+)
+from sglang.test.test_utils import CustomTestCase
+
+NUM_EXPERTS = 4
+PREFIX = "model.layers.0.block_sparse_moe"
+
+
+class _RecordingParam:
+    """Deliberately rejects `return_success`: no loader in the tree accepts
+    it, so a caller passing it raises TypeError here."""
+
+    def __init__(self, calls):
+        self._calls = calls
+
+    def weight_loader(self, param, loaded_weight, name, shard_id, expert_id):
+        self._calls.append(
+            {
+                "name": name,
+                "shard_id": shard_id,
+                "expert_id": expert_id,
+                "value": loaded_weight,
+            }
+        )
+
+
+def _mapping(num_experts=NUM_EXPERTS):
+    return FusedMoE.make_expert_params_mapping(
+        ckpt_gate_proj_name="gate_proj",
+        ckpt_down_proj_name="down_proj",
+        ckpt_up_proj_name="up_proj",
+        num_experts=num_experts,
+    )
+
+
+def _run(weights, params_dict=None, calls=None):
+    """Drive the loader, returning (passed_through, calls)."""
+    calls = [] if calls is None else calls
+    if params_dict is None:
+        params_dict = _AutoParams(calls)
+    passed_through = list(
+        granitemoe_load_split_experts(
+            weights,
+            expert_params_mapping=_mapping(),
+            params_dict=params_dict,
+        )
+    )
+    return passed_through, calls
+
+
+class _AutoParams(dict):
+    """params_dict that materialises a recording param for any requested name."""
+
+    def __init__(self, calls):
+        super().__init__()
+        self._calls = calls
+
+    def __getitem__(self, key):
+        if key not in self:
+            super().__setitem__(key, _RecordingParam(self._calls))
+        return super().__getitem__(key)
+
+
+class TestGraniteMoeLoadSplitExperts(CustomTestCase):
+    def test_split_experts_are_loaded_not_passed_through(self):
+        """The bug: these tensors reached the generic path and were dropped."""
+        weights = [
+            (f"{PREFIX}.experts.3.gate_proj.weight", torch.zeros(2)),
+            (f"{PREFIX}.experts.3.up_proj.weight", torch.zeros(2)),
+            (f"{PREFIX}.experts.1.down_proj.weight", torch.zeros(2)),
+        ]
+        passed_through, calls = _run(weights)
+
+        self.assertEqual(passed_through, [], "split experts must not fall through")
+        self.assertEqual(len(calls), 3)
+        self.assertEqual(
+            [(c["name"], c["shard_id"], c["expert_id"]) for c in calls],
+            [
+                (f"{PREFIX}.experts.w13_weight", "w1", 3),
+                (f"{PREFIX}.experts.w13_weight", "w3", 3),
+                (f"{PREFIX}.experts.w2_weight", "w2", 1),
+            ],
+        )
+
+    def test_scale_keeps_its_suffix(self):
+        """FusedMoE's loader dispatches on substrings of the name it is handed,
+        so a scale must arrive as `*_weight_scale`. Handing it `*_weight` would
+        load the scale as if it were the weight."""
+        weights = [
+            (f"{PREFIX}.experts.0.gate_proj.weight_scale", torch.zeros(1)),
+            (f"{PREFIX}.experts.0.down_proj.weight_scale", torch.zeros(1)),
+        ]
+        _, calls = _run(weights)
+
+        self.assertEqual(
+            [c["name"] for c in calls],
+            [
+                f"{PREFIX}.experts.w13_weight_scale",
+                f"{PREFIX}.experts.w2_weight_scale",
+            ],
+        )
+
+    def test_every_expert_and_projection_is_loaded(self):
+        """A missing (expert, projection) pair is the silent-drop bug: nothing
+        matches and those weights never reach the layer."""
+        weights = [
+            (f"{PREFIX}.experts.{e}.{proj}.{suffix}", torch.zeros(1))
+            for e in range(NUM_EXPERTS)
+            for proj in ("gate_proj", "up_proj", "down_proj")
+            for suffix in ("weight", "weight_scale")
+        ]
+        passed_through, calls = _run(weights)
+
+        self.assertEqual(passed_through, [])
+        self.assertEqual(len(calls), len(weights))
+        self.assertEqual(
+            {(c["expert_id"], c["shard_id"]) for c in calls},
+            {(e, s) for e in range(NUM_EXPERTS) for s in ("w1", "w2", "w3")},
+        )
+
+    def test_non_expert_tensors_pass_through_untouched(self):
+        weights = [
+            ("model.layers.0.self_attn.q_proj.weight", torch.zeros(1)),
+            ("model.layers.0.self_attn.q_proj.weight_scale", torch.zeros(1)),
+            (f"{PREFIX}.router.layer.weight", torch.zeros(1)),
+            ("model.layers.0.input_layernorm.weight", torch.zeros(1)),
+            ("lm_head.weight", torch.zeros(1)),
+        ]
+        passed_through, calls = _run(weights)
+
+        self.assertEqual([n for n, _ in passed_through], [n for n, _ in weights])
+        self.assertEqual(calls, [])
+
+    def test_packed_layout_experts_pass_through(self):
+        """The packed (unquantized) path pre-splits into w1/w2/w3, which load
+        downstream. Claiming or rejecting them here breaks the bf16 model."""
+        weights = [
+            (f"{PREFIX}.experts.0.{shard}.weight", torch.zeros(1))
+            for shard in ("w1", "w2", "w3")
+        ]
+        passed_through, calls = _run(weights)
+
+        self.assertEqual([n for n, _ in passed_through], [n for n, _ in weights])
+        self.assertEqual(calls, [])
+
+    def test_unrecognised_expert_tensor_raises(self):
+        """The original failure was silent. An expert tensor that matches no
+        mapping must now fail the load instead of producing a garbage model."""
+        weights = [(f"{PREFIX}.experts.0.mystery_proj.weight", torch.zeros(1))]
+        with self.assertRaises(ValueError) as ctx:
+            _run(weights)
+        self.assertIn("unmatched MoE expert tensor", str(ctx.exception))
+
+
+class TestIsPackedExpert(CustomTestCase):
+    """A false positive here silently reintroduces the dropped-expert bug:
+    an unmatched expert tensor passes through instead of raising."""
+
+    def test_shard_substring_requires_dot_delimiters(self):
+        """`w1` appearing inside a longer segment is not a packed shard. Without
+        the dots this predicate would swallow such a tensor as packed and the
+        loader would drop it silently instead of raising."""
+        for name in (
+            f"{PREFIX}.experts.0.w1_proj.weight",
+            f"{PREFIX}.experts.0.gate_w2.weight",
+            f"{PREFIX}.experts.0.w3x.weight",
+        ):
+            with self.subTest(name=name):
+                self.assertFalse(_is_packed_expert(name))
+
+
+if __name__ == "__main__":
+    unittest.main()