Fix Muse Glimmer ModelOpt mixed weight mapping (#37510)

python/sglang/srt/models/muse_glimmer.py

@@ -13,7 +13,6 @@
 # ==============================================================================
 
 import logging
-import re
 from typing import Iterable, List, Optional, Tuple
 
 import torch
@@ -57,7 +56,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
-from sglang.srt.models.utils import apply_qk_norm, permute_inv
+from sglang.srt.models.utils import WeightsMapper, apply_qk_norm, permute_inv
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_cuda
 
@@ -84,22 +83,26 @@ _VISION_NAME_FRAGMENTS = (
     "perception_emb_norm",
 )
 
-# Vendor tensor names -> this port's; applied simultaneously.
+# Shared by _vendor_weight_name and hf_to_sglang_mapper;
-_VENDOR_RENAMES = {
+# a rule missing from one silently breaks the other.
-    "post_attention_layernorm": "post_attn_norm",
+_VENDOR_TO_SGLANG = WeightsMapper(
-    "pre_feedforward_layernorm": "post_attention_layernorm",
+    orig_to_new_prefix={
-    "post_feedforward_layernorm": "post_ffn_norm",
+        "model.language_model.": "model.",
-    "self_attn.gate_proj": "self_attn.output_gate_proj",
+        # The vision modules hang off the entry class, not off ``model``.
-}
+        "model.vision_": "vision_",
-
+    },
-_VENDOR_RENAME_RE = re.compile("|".join(re.escape(key) for key in _VENDOR_RENAMES))
+    # Only the first matching substring is applied; keep these non-overlapping.
+    orig_to_new_substr={
+        "post_attention_layernorm": "post_attn_norm",
+        "pre_feedforward_layernorm": "post_attention_layernorm",
+        "post_feedforward_layernorm": "post_ffn_norm",
+        "self_attn.gate_proj": "self_attn.output_gate_proj",
+    },
+)
 
 
 def _vendor_weight_name(name: str) -> str:
-    name = name.replace("model.language_model.", "model.", 1)
+    return _VENDOR_TO_SGLANG.apply_list([name])[0]
-    # The vision modules hang off the entry class, not off ``model``.
-    name = name.replace("model.vision_", "vision_", 1)
-    return _VENDOR_RENAME_RE.sub(lambda m: _VENDOR_RENAMES[m.group(0)], name)
 
 
 def get_attention_sliding_window_size(config) -> int:
@@ -928,6 +931,8 @@ class MuseGlimmerForCausalLM(nn.Module):
 class MuseGlimmerForConditionalGeneration(MuseGlimmerForCausalLM):
     """Vendor multimodal HF export: the MuseGlimmerForCausalLM decoder plus the image tower."""
 
+    # Only this class reads vendor-named checkpoints, so only it needs the mapper.
+    hf_to_sglang_mapper = _VENDOR_TO_SGLANG
     checkpoint_uses_vendor_names = True
     builds_vision_tower = True
 

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -40,6 +40,7 @@ from sglang.srt.model_loader.weight_utils import (
     get_quant_config,
 )
 from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
+from sglang.srt.models.muse_glimmer import MuseGlimmerForConditionalGeneration
 from sglang.srt.models.utils import WeightsMapper
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -889,6 +890,69 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase):
             ["language_model.lm_head", "lm_head"],
         )
 
+    def test_muse_glimmer_mixed_precision_resolves_runtime_names(self):
+        """The vendor keys quant metadata under ``model.language_model.*``;
+        it must resolve for the ``model.*`` modules the runtime builds.
+        """
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "model.language_model.layers.0.mlp.gate_proj": {
+                        "quant_algo": "W4A16_NVFP4",
+                        "group_size": 16,
+                    },
+                    "model.language_model.layers.0.mlp.up_proj": {
+                        "quant_algo": "W4A16_NVFP4",
+                        "group_size": 16,
+                    },
+                    "model.language_model.layers.0.self_attn.q_proj": {
+                        "quant_algo": "FP8"
+                    },
+                    "model.language_model.layers.0.self_attn.k_proj": {
+                        "quant_algo": "FP8"
+                    },
+                    "model.language_model.layers.0.self_attn.v_proj": {
+                        "quant_algo": "FP8"
+                    },
+                    "model.language_model.layers.0.self_attn.gate_proj": {
+                        "quant_algo": "FP8"
+                    },
+                    "lm_head": {"quant_algo": "W4A16_NVFP4", "group_size": 16},
+                    "model.vision_tower.layers.0.attn.q_proj": {"quant_algo": "FP8"},
+                },
+                "packed_modules_mapping": (
+                    MuseGlimmerForConditionalGeneration.packed_modules_mapping
+                ),
+            }
+        )
+        quant_config.apply_weight_name_mapper(
+            MuseGlimmerForConditionalGeneration.hf_to_sglang_mapper
+        )
+
+        self.assertEqual(
+            quant_config._resolve_quant_algo("model.layers.0.mlp.gate_up_proj"),
+            "W4A16_NVFP4",
+        )
+        # Attention stays unfused whenever a quant_config is present, so q/k/v
+        # resolve per shard; only the MLP goes through packed_modules_mapping.
+        self.assertEqual(
+            quant_config._resolve_quant_algo("model.layers.0.self_attn.q_proj"),
+            "FP8",
+        )
+        self.assertEqual(
+            quant_config._resolve_quant_algo(
+                "model.layers.0.self_attn.output_gate_proj"
+            ),
+            "FP8",
+        )
+        self.assertEqual(quant_config._resolve_quant_algo("lm_head"), "W4A16_NVFP4")
+        # The vision tower hangs off the entry class, not off ``model``.
+        self.assertEqual(
+            quant_config._resolve_quant_algo("vision_tower.layers.0.attn.q_proj"),
+            "FP8",
+        )
+
     def test_nemotron_mixed_precision_with_nvfp4_layers_uses_modelopt_mixed(self):
         model_config = ModelConfig.__new__(ModelConfig)
         model_config.hf_config = MagicMock()