diff --git a/python/sglang/srt/models/muse_glimmer.py b/python/sglang/srt/models/muse_glimmer.py index 07d2599f3..918dd5194 100644 --- a/python/sglang/srt/models/muse_glimmer.py +++ b/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. -_VENDOR_RENAMES = { - "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", -} - -_VENDOR_RENAME_RE = re.compile("|".join(re.escape(key) for key in _VENDOR_RENAMES)) +# Shared by _vendor_weight_name and hf_to_sglang_mapper; +# a rule missing from one silently breaks the other. +_VENDOR_TO_SGLANG = WeightsMapper( + orig_to_new_prefix={ + "model.language_model.": "model.", + # The vision modules hang off the entry class, not off ``model``. + "model.vision_": "vision_", + }, + # 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) - # 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) + return _VENDOR_TO_SGLANG.apply_list([name])[0] 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 diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index fa607a540..80230b90b 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/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()