From ae1acf822dd357641d885f30c58d2ad547ec9220 Mon Sep 17 00:00:00 2001 From: jlqibm <87678480+jlqibm@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:10:20 -0400 Subject: [PATCH] [GraniteMoE] Load split per-expert quantized MoE weights (#37679) Co-authored-by: Claude Opus 4.8 Co-authored-by: Xinyuan Tong Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> --- python/sglang/srt/models/granitemoe.py | 62 ++++++ .../test_granitemoe_split_expert_loading.py | 195 ++++++++++++++++++ 2 files changed, 257 insertions(+) create mode 100644 test/registered/unit/models/test_granitemoe_split_expert_loading.py diff --git a/python/sglang/srt/models/granitemoe.py b/python/sglang/srt/models/granitemoe.py index b38814ae9..7ee66789a 100644 --- a/python/sglang/srt/models/granitemoe.py +++ b/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) diff --git a/test/registered/unit/models/test_granitemoe_split_expert_loading.py b/test/registered/unit/models/test_granitemoe_split_expert_loading.py new file mode 100644 index 000000000..730b40d4b --- /dev/null +++ b/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..{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()