From f7a404e9c308502ab87f3acb625052d261f73a99 Mon Sep 17 00:00:00 2001 From: vikram singh shekhawat Date: Mon, 17 Aug 2026 12:23:59 +0530 Subject: [PATCH] Fix rope config compatibility and VL/transformers-fallback weight loading (#31575) Co-authored-by: Claude Sonnet 4.5 (1M context) --- python/sglang/srt/managers/mm_schedule.py | 7 +- python/sglang/srt/models/ernie45_moe_vl.py | 4 +- python/sglang/srt/models/olmo2.py | 3 +- python/sglang/srt/models/qwen.py | 8 +- python/sglang/srt/models/transformers.py | 47 +++++++ .../srt/utils/hf_transformers/common.py | 3 +- .../test_transformers_collect_mm_kwargs.py | 126 ++++++++++++++++++ 7 files changed, 190 insertions(+), 8 deletions(-) create mode 100644 test/manual/models/test_transformers_collect_mm_kwargs.py diff --git a/python/sglang/srt/managers/mm_schedule.py b/python/sglang/srt/managers/mm_schedule.py index ed291525f..45dd91bbe 100644 --- a/python/sglang/srt/managers/mm_schedule.py +++ b/python/sglang/srt/managers/mm_schedule.py @@ -9,12 +9,13 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.multimodal.evs import EVSEmbeddingResult from sglang.srt.runtime_context import get_parallel, get_schedule -from sglang.srt.utils import is_hip, is_npu +from sglang.srt.utils import is_hip, is_npu, is_xpu from sglang.srt.utils.async_probe import maybe_assert_sum from sglang.utils import logger _is_hip = is_hip() _is_npu = is_npu() +_is_xpu = is_xpu() embedding_cache: Optional[MultiModalStaticCache] = None @@ -510,9 +511,9 @@ def _get_chunked_prefill_embedding( is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req) if is_per_image: - if _is_hip or _is_npu: + if _is_hip or _is_npu or _is_xpu: # ROCm CI regressed with one large cross-request ViT batch; keep - # the previous per-request path on HIP while CUDA uses batching. + # the previous per-request path on HIP/NPU/XPU while CUDA uses batching. chunk = _get_chunked_embedding_by_item( data_embedding_func, embedding_items_per_req, diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index 7321a59db..0879d3184 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/python/sglang/srt/models/ernie45_moe_vl.py @@ -44,6 +44,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils.hf_transformers.common import get_rope_config logger = logging.getLogger(__name__) @@ -368,8 +369,7 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module): prefix: str = "", ): super().__init__() - rope_theta = config.rope_parameters["rope_theta"] - rope_scaling = config.rope_parameters + rope_theta, rope_scaling = get_rope_config(config) rope_is_neox_style = getattr(config, "rope_is_neox_style", False) freq_allocation = getattr(config, "freq_allocation", 20) max_position_embeddings = getattr(config, "max_position_embeddings", 131072) diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index a743cd75a..85d5f507e 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/python/sglang/srt/models/olmo2.py @@ -49,6 +49,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import add_prefix, is_cuda, make_layers +from sglang.srt.utils.hf_transformers.common import get_rope_config _is_cuda = is_cuda() @@ -100,7 +101,7 @@ class Olmo2Attention(nn.Module): self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.max_position_embeddings = config.max_position_embeddings - self.rope_theta = config.rope_parameters["rope_theta"] + self.rope_theta, _ = get_rope_config(config) # Attention input projection. Projects x -> (q, k, v) self.qkv_proj = QKVParallelLinear( diff --git a/python/sglang/srt/models/qwen.py b/python/sglang/srt/models/qwen.py index 66fdf1b09..f6fd92128 100644 --- a/python/sglang/srt/models/qwen.py +++ b/python/sglang/srt/models/qwen.py @@ -337,7 +337,10 @@ class QWenLMHeadModel(nn.Module): for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue - name = name.replace(weight_name, param_name) + temp_name = name.replace(weight_name, param_name) + if temp_name not in params_dict: + continue + name = temp_name # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue @@ -349,6 +352,9 @@ class QWenLMHeadModel(nn.Module): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue + # Skip visual encoder weights (e.g. Qwen-VL-Chat transformer.visual.*) + if name not in params_dict: + continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) diff --git a/python/sglang/srt/models/transformers.py b/python/sglang/srt/models/transformers.py index a5584c85e..20d600e90 100644 --- a/python/sglang/srt/models/transformers.py +++ b/python/sglang/srt/models/transformers.py @@ -1345,6 +1345,22 @@ class MultiModalMixin: super().__init__(*args, **kwargs) self._mm_padding_pattern = MultiModalityDataPaddingPatternMultimodalTokens() + # transformers v5 flattened SigLIP/CLIP (dropped the "vision_model" + # wrapper); older checkpoints still ship "vision_tower.vision_model.*" + # keys, so remap them when the live model lacks that sub-module. + vt = getattr(self.model, "vision_tower", None) + if vt is not None and not any( + name == "vision_model" for name, _ in vt.named_children() + ): + self.weight_mapper = ( + WeightsMapper( + orig_to_new_prefix={ + "vision_tower.vision_model.": "model.vision_tower.", + } + ) + | self.weight_mapper + ) + def _uses_mrope_positions(self) -> bool: rope_scaling = getattr(self.text_config, "rope_scaling", None) if isinstance(rope_scaling, Mapping) and "mrope_section" in rope_scaling: @@ -1497,6 +1513,14 @@ class MultiModalMixin: ): mm_inputs = forward_batch.mm_inputs target_device = next(self.model.parameters()).device + # 5D features (num_images, num_patches, C, H, W) can't be flattened + # here: anyres models pad num_patches per HF processor call, so a + # flattened concat would leave stray padding rows once items with + # different tile counts are combined. Defer them and pad to the + # batch-wide max instead -- the model's own get_image_features + # re-derives each image's real patch count from image_sizes and + # slices the padding back out. + pending_5d_features: dict = {} for batch_idx in range(len(mm_inputs or [])): mm_input = mm_inputs[batch_idx] @@ -1519,6 +1543,11 @@ class MultiModalMixin: feature = item.feature if isinstance(feature, torch.Tensor): feature = feature.to(device=target_device) + if feature.dim() == 5: + pending_5d_features.setdefault(feature_key, []).append( + feature + ) + continue if feature_key not in kwargs: kwargs[feature_key] = feature elif isinstance(feature, torch.Tensor) and isinstance( @@ -1528,6 +1557,24 @@ class MultiModalMixin: [kwargs[feature_key], feature], dim=0 ) + for feature_key, tensors in pending_5d_features.items(): + max_patches = max(t.shape[1] for t in tensors) + padded = [] + for t in tensors: + if t.shape[1] < max_patches: + pad = t.new_zeros( + (t.shape[0], max_patches - t.shape[1], *t.shape[2:]) + ) + t = torch.cat([t, pad], dim=1) + padded.append(t) + combined = torch.cat(padded, dim=0) + if feature_key in kwargs: + kwargs[feature_key] = torch.cat( + [kwargs[feature_key], combined], dim=0 + ) + else: + kwargs[feature_key] = combined + return kwargs def _forward_hidden_states( diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index efc8dee39..de126595c 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -365,7 +365,8 @@ def get_rope_config(config): """ rope_params = getattr(config, "rope_parameters", None) if rope_params is not None: - return rope_params["rope_theta"], rope_params + rope_theta = rope_params.get("rope_theta", getattr(config, "rope_theta", 10000)) + return rope_theta, rope_params return getattr(config, "rope_theta", 10000), getattr(config, "rope_scaling", None) diff --git a/test/manual/models/test_transformers_collect_mm_kwargs.py b/test/manual/models/test_transformers_collect_mm_kwargs.py new file mode 100644 index 000000000..e2ef629b6 --- /dev/null +++ b/test/manual/models/test_transformers_collect_mm_kwargs.py @@ -0,0 +1,126 @@ +""" +Unit tests for MultiModalMixin._collect_mm_kwargs' handling of 5D +pixel_values features in the generic Transformers fallback backend +(sglang.srt.models.transformers). +""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.models.transformers import MultiModalMixin + + +def _make_item(modality_name, feature, model_specific_data=None): + return SimpleNamespace( + modality=SimpleNamespace(name=modality_name), + feature=feature, + model_specific_data=model_specific_data or {}, + ) + + +def _make_mm_input(items): + return SimpleNamespace(mm_items=items) + + +def _make_forward_batch(mm_inputs, is_decode=False, contains_mm_inputs=True): + return SimpleNamespace( + token_type_ids=None, + forward_mode=SimpleNamespace(is_decode=lambda: is_decode), + mm_inputs=mm_inputs, + contains_mm_inputs=lambda: contains_mm_inputs, + ) + + +def _make_self(): + """Lightweight stand-in for a TransformersForCausalLM instance -- only + `.model` (for the device lookup) and the mixin's own feature-key map + are actually used by `_collect_mm_kwargs`.""" + return SimpleNamespace( + model=torch.nn.Linear(1, 1), + _mm_feature_kwarg=MultiModalMixin._mm_feature_kwarg, + ) + + +class TestCollectMmKwargs5DPadding(unittest.TestCase): + def test_equal_patch_counts_no_padding(self): + """Sanity check: same num_patches across items concatenates cleanly.""" + item1 = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0)) + item2 = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 2.0)) + forward_batch = _make_forward_batch( + [_make_mm_input([item1]), _make_mm_input([item2])] + ) + + kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch) + + pixel_values = kwargs["pixel_values"] + self.assertEqual(pixel_values.shape, (2, 3, 3, 4, 4)) + self.assertTrue(torch.all(pixel_values[0] == 1.0)) + self.assertTrue(torch.all(pixel_values[1] == 2.0)) + + def test_different_patch_counts_padded_to_batch_max(self): + """Test: items with a different tile count must be + zero-padded to the batch-wide max num_patches, not just concatenated + as-is (which would crash on mismatched shapes or misalign data).""" + item_small = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0)) # 3 patches + item_large = _make_item("IMAGE", torch.full((1, 5, 3, 4, 4), 2.0)) # 5 patches + forward_batch = _make_forward_batch( + [_make_mm_input([item_small]), _make_mm_input([item_large])] + ) + + kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch) + + pixel_values = kwargs["pixel_values"] + self.assertEqual(pixel_values.shape, (2, 5, 3, 4, 4)) + # item_small's real 3 patches are preserved... + self.assertTrue(torch.all(pixel_values[0, :3] == 1.0)) + # ...and its padding (patches 3-4) is zeroed, not garbage/leftover data. + self.assertTrue(torch.all(pixel_values[0, 3:] == 0.0)) + # item_large needed no padding at all. + self.assertTrue(torch.all(pixel_values[1] == 2.0)) + + def test_multi_image_item_with_different_patch_counts_within_one_item(self): + """A single multi-image item/request can itself already contain + per-image padding applied by the HF processor; the batch-level + padding must still pad up to the overall max without disturbing it.""" + # 2 images already padded to 4 patches by the HF processor, batched + # against another item that only needed 2 patches. + item_multi_image = _make_item("IMAGE", torch.full((2, 4, 3, 4, 4), 1.0)) + item_single = _make_item("IMAGE", torch.full((1, 2, 3, 4, 4), 2.0)) + forward_batch = _make_forward_batch( + [_make_mm_input([item_multi_image]), _make_mm_input([item_single])] + ) + + kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch) + + pixel_values = kwargs["pixel_values"] + self.assertEqual(pixel_values.shape, (3, 4, 3, 4, 4)) + self.assertTrue(torch.all(pixel_values[:2] == 1.0)) + self.assertTrue(torch.all(pixel_values[2, :2] == 2.0)) + self.assertTrue(torch.all(pixel_values[2, 2:] == 0.0)) + + def test_decode_mode_skips_collection(self): + """During decode (no new mm inputs to process this step), no + multimodal kwargs should be produced even if mm_inputs is present.""" + item = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0)) + forward_batch = _make_forward_batch([_make_mm_input([item])], is_decode=True) + + kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch) + + self.assertNotIn("pixel_values", kwargs) + + def test_non_image_modality_uses_correct_feature_key(self): + """Video features (also potentially 5D) must land under their own + kwarg key, not be mixed in with image pixel_values.""" + item = _make_item("VIDEO", torch.full((1, 3, 3, 4, 4), 1.0)) + forward_batch = _make_forward_batch([_make_mm_input([item])]) + + kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch) + + self.assertIn("pixel_values_videos", kwargs) + self.assertNotIn("pixel_values", kwargs) + + +if __name__ == "__main__": + unittest.main()