[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>

python/sglang/srt/models/granitemoe.py

@@ -54,6 +54,52 @@ def granitemoe_split_expert_weights(
             yield name, weight
 
 
+def _match_split_expert(
+    name: str,
+    expert_params_mapping: list[tuple[str, str, int, str]],
+) -> Optional[tuple[str, int, str]]:
+    """Resolve a split-expert tensor name to (param_name, expert_id, shard_id)."""
+    for param_name, weight_name, expert_id, shard_id in expert_params_mapping:
+        if weight_name in name:
+            # Keeping the checkpoint's trailing suffix (`.weight` vs
+            # `.weight_scale`) routes scales to FusedMoE's scale branch.
+            return name.replace(weight_name, param_name), expert_id, shard_id
+    return None
+
+
+def _is_packed_expert(name: str) -> bool:
+    """Whether an expert tensor is the packed layout, already split to w1/w2/w3."""
+    return any(f".{shard}." in name for shard in ("w1", "w2", "w3"))
+
+
+def granitemoe_load_split_experts(
+    weights: Iterable[tuple[str, torch.Tensor]],
+    expert_params_mapping: list[tuple[str, str, int, str]],
+    params_dict: dict[str, torch.nn.Parameter],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Load split quantized experts, yielding the tensors it does not claim."""
+    for name, loaded_weight in weights:
+        match = _match_split_expert(name, expert_params_mapping)
+
+        if match is None:
+            # Packed-layout experts load downstream; any other expert tensor
+            # would be dropped silently and generate garbage, so fail loudly.
+            if ".block_sparse_moe.experts." in name and not _is_packed_expert(name):
+                raise ValueError(f"unmatched MoE expert tensor: {name}")
+            yield name, loaded_weight
+            continue
+
+        mapped_name, expert_id, shard_id = match
+        param = params_dict[mapped_name]
+        param.weight_loader(
+            param,
+            loaded_weight,
+            mapped_name,
+            shard_id=shard_id,
+            expert_id=expert_id,
+        )
+
+
 class GraniteMoeMoE(nn.Module):
     """A tensor-parallel MoE implementation for GraniteMoe that shards each
     expert across all ranks.
@@ -443,10 +489,26 @@ class GraniteMoeForCausalLM(nn.Module):
         else:
             return self.pooler(hidden_states, forward_batch)
 
+    def _split_expert_params_mapping(self) -> list[tuple[str, str, int, str]]:
+        # llmcompressor emits one tensor per expert per projection
+        # (`experts.{e}.gate_proj.weight`, `.weight_scale`, ...), not the packed
+        # input_linear/output_linear that granitemoe_split_expert_weights maps.
+        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=self.config.num_local_experts,
+        )
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = granitemoe_split_expert_weights(
             self.hf_to_sglang_mapper.apply(weights)
         )
+        weights = granitemoe_load_split_experts(
+            weights,
+            expert_params_mapping=self._split_expert_params_mapping(),
+            params_dict=dict(self.named_parameters()),
+        )
         mixtral.MixtralForCausalLM.load_weights(self, weights)
 
 

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()