[Bugfix][Model] Fix M-RoPE offset double-count in Qwen3-Omni (#58890)
Signed-off-by: C1ves <3121985758@qq.com>
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.models.qwen3_omni_moe_thinker import (
|
||||
Qwen3OmniMoeThinkerForConditionalGeneration,
|
||||
)
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalFieldElem,
|
||||
MultiModalKwargsItem,
|
||||
PlaceholderRange,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def _force_cpu_default_device():
|
||||
# get_mrope_input_positions returns CPU tensors (via torch.from_numpy).
|
||||
# Ensure the default device is CPU so the rest of the test tensors match.
|
||||
original = torch.get_default_device()
|
||||
torch.set_default_device("cpu")
|
||||
yield
|
||||
torch.set_default_device(original)
|
||||
|
||||
|
||||
TEXT_TOKEN_ID = 1
|
||||
VISION_START_TOKEN_ID = 100
|
||||
VISION_END_TOKEN_ID = 101
|
||||
IMAGE_TOKEN_ID = 200
|
||||
VIDEO_TOKEN_ID = 201
|
||||
AUDIO_TOKEN_ID = 202
|
||||
AUDIO_START_TOKEN_ID = 300
|
||||
AUDIO_END_TOKEN_ID = 301
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyVisionConfig:
|
||||
spatial_merge_size: int = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyConfig:
|
||||
position_id_per_seconds: int = 25
|
||||
vision_config: DummyVisionConfig = dataclasses.field(
|
||||
default_factory=DummyVisionConfig
|
||||
)
|
||||
|
||||
|
||||
def make_model() -> Qwen3OmniMoeThinkerForConditionalGeneration:
|
||||
model = Qwen3OmniMoeThinkerForConditionalGeneration.__new__(
|
||||
Qwen3OmniMoeThinkerForConditionalGeneration
|
||||
)
|
||||
model.config = DummyConfig()
|
||||
return model
|
||||
|
||||
|
||||
def _elem(data):
|
||||
return MultiModalFieldElem(
|
||||
data=torch.tensor(data),
|
||||
field=None, # HACK.
|
||||
)
|
||||
|
||||
|
||||
def make_image_feature(offset: int, grid_thw: tuple[int, int, int]):
|
||||
return MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem({"image_grid_thw": _elem(grid_thw)}),
|
||||
modality="image",
|
||||
identifier="DUMMY",
|
||||
mm_position=PlaceholderRange(
|
||||
offset=offset, length=int(torch.tensor(grid_thw).prod())
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_video_feature(
|
||||
offset: int, grid_thw: tuple[int, int, int], use_audio_in_video: bool = False
|
||||
):
|
||||
data = {"video_grid_thw": _elem(grid_thw)}
|
||||
if use_audio_in_video:
|
||||
data["use_audio_in_video"] = _elem(True)
|
||||
return MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem(data),
|
||||
modality="video",
|
||||
identifier="DUMMY",
|
||||
mm_position=PlaceholderRange(
|
||||
offset=offset, length=int(torch.tensor(grid_thw).prod())
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_audio_feature(offset: int, audio_feature_length: int):
|
||||
return MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem(
|
||||
{"audio_feature_lengths": _elem(audio_feature_length)}
|
||||
),
|
||||
modality="audio",
|
||||
identifier="DUMMY",
|
||||
mm_position=PlaceholderRange(offset=offset, length=0),
|
||||
)
|
||||
|
||||
|
||||
def test_mrope_positions_image_at_prompt_end():
|
||||
"""A prompt ending with <|vision_end|> must still produce seq_len positions."""
|
||||
model = make_model()
|
||||
# [<|vision_start|>, <|image_pad|> x 4, <|vision_end|>]
|
||||
input_tokens = (
|
||||
[VISION_START_TOKEN_ID] + [IMAGE_TOKEN_ID] * 4 + [VISION_END_TOKEN_ID]
|
||||
)
|
||||
mm_features = [make_image_feature(offset=1, grid_thw=(1, 2, 2))]
|
||||
|
||||
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
|
||||
|
||||
def test_mrope_positions_audio_at_prompt_end():
|
||||
model = make_model()
|
||||
# audio_feature_length=16 -> 2 audio tokens.
|
||||
# [<|audio_start|>, <|audio_pad|> x 2, <|audio_end|>]
|
||||
input_tokens = [AUDIO_START_TOKEN_ID] + [AUDIO_TOKEN_ID] * 2 + [AUDIO_END_TOKEN_ID]
|
||||
mm_features = [make_audio_feature(offset=1, audio_feature_length=16)]
|
||||
|
||||
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
|
||||
|
||||
def test_mrope_positions_video_at_prompt_end():
|
||||
model = make_model()
|
||||
# [T, <|vision_start|>, <|video_pad|> x 8, <|vision_end|>]
|
||||
input_tokens = (
|
||||
[TEXT_TOKEN_ID]
|
||||
+ [VISION_START_TOKEN_ID]
|
||||
+ [VIDEO_TOKEN_ID] * 8
|
||||
+ [VISION_END_TOKEN_ID]
|
||||
)
|
||||
mm_features = [make_video_feature(offset=2, grid_thw=(2, 2, 2))]
|
||||
|
||||
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
|
||||
|
||||
def test_mrope_positions_audio_in_video_at_prompt_end():
|
||||
model = make_model()
|
||||
# video grid (1, 2, 2) -> 4 video tokens; audio_feature_length=16 -> 2
|
||||
# audio tokens, interleaved after the video tokens.
|
||||
# [T, T, <|vision_start|>, <|audio_start|>, V x 4, A x 2,
|
||||
# <|audio_end|>, <|vision_end|>]
|
||||
input_tokens = (
|
||||
[TEXT_TOKEN_ID] * 2
|
||||
+ [VISION_START_TOKEN_ID, AUDIO_START_TOKEN_ID]
|
||||
+ [VIDEO_TOKEN_ID] * 4
|
||||
+ [AUDIO_TOKEN_ID] * 2
|
||||
+ [AUDIO_END_TOKEN_ID, VISION_END_TOKEN_ID]
|
||||
)
|
||||
mm_features = [
|
||||
make_video_feature(offset=3, grid_thw=(1, 2, 2), use_audio_in_video=True),
|
||||
make_audio_feature(offset=3, audio_feature_length=16),
|
||||
]
|
||||
|
||||
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
|
||||
|
||||
def test_mrope_positions_image_then_text():
|
||||
"""Media followed by text: positions must match the expected layout exactly.
|
||||
|
||||
Before the fix this case passed the length guard but every position from
|
||||
the first <|image_pad|> onward was silently shifted by one slot.
|
||||
"""
|
||||
model = make_model()
|
||||
# [T, T, <|vision_start|>, <|image_pad|> x 4, <|vision_end|>, T, T, T]
|
||||
input_tokens = (
|
||||
[TEXT_TOKEN_ID] * 2
|
||||
+ [VISION_START_TOKEN_ID]
|
||||
+ [IMAGE_TOKEN_ID] * 4
|
||||
+ [VISION_END_TOKEN_ID]
|
||||
+ [TEXT_TOKEN_ID] * 3
|
||||
)
|
||||
mm_features = [make_image_feature(offset=3, grid_thw=(1, 2, 2))]
|
||||
|
||||
positions, delta = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
expected = torch.tensor(
|
||||
[
|
||||
[0, 1, 2, 3, 3, 3, 3, 5, 6, 7, 8],
|
||||
[0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 8],
|
||||
[0, 1, 2, 3, 4, 3, 4, 5, 6, 7, 8],
|
||||
]
|
||||
)
|
||||
torch.testing.assert_close(positions, expected)
|
||||
assert delta == 9 - len(input_tokens)
|
||||
@@ -2260,10 +2260,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
)
|
||||
st_idx += text_len
|
||||
|
||||
bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(bos_pos)
|
||||
st_idx += 1
|
||||
|
||||
if modality == "audio":
|
||||
audio_tokens = self._compute_audio_token_count(
|
||||
data["audio_feature_length"]
|
||||
@@ -2276,7 +2272,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
|
||||
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
st = offset + 1 + audio_tokens + 1
|
||||
st = offset + audio_tokens + 1
|
||||
|
||||
elif modality == "image":
|
||||
grid_t = data["grid_t"]
|
||||
@@ -2294,7 +2290,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
|
||||
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
st = offset + 1 + image_len + 1
|
||||
st = offset + image_len + 1
|
||||
|
||||
elif modality == "video":
|
||||
grid_t = data["grid_t"]
|
||||
@@ -2313,7 +2309,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
|
||||
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
st = offset + 1 + video_len + 1
|
||||
st = offset + video_len + 1
|
||||
else:
|
||||
audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
|
||||
llm_pos_ids_list.append(audio_bos_pos)
|
||||
@@ -2330,7 +2326,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
audio_len = self._compute_audio_token_count(
|
||||
data["audio_feature_length"]
|
||||
)
|
||||
st = offset + 2 + video_len + audio_len + 2
|
||||
st = offset + 1 + video_len + audio_len + 2
|
||||
|
||||
if st < seq_len:
|
||||
st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
|
||||
|
||||
Reference in New Issue
Block a user