[Bugfix][Model] Fix Qwen3-Omni interleaved M-RoPE boundary positions (#59842)
Signed-off-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
co-authored by
OpenAI Codex
parent
48eb44b719
commit
8c417a2dae
@@ -145,27 +145,45 @@ def test_mrope_positions_video_at_prompt_end():
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
|
||||
|
||||
def test_mrope_positions_audio_in_video_at_prompt_end():
|
||||
@pytest.mark.parametrize("num_videos", [1, 2])
|
||||
@pytest.mark.parametrize("trailing_tokens", [0, 3])
|
||||
def test_mrope_positions_audio_in_video_boundaries(num_videos, trailing_tokens):
|
||||
"""Interleaved BOS/EOS positions and decode delta match Qwen3-Omni HF."""
|
||||
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_tokens = (
|
||||
[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),
|
||||
]
|
||||
input_tokens = [TEXT_TOKEN_ID] * 2
|
||||
mm_features = []
|
||||
expected = [torch.tensor([[0, 1]]).expand(3, -1)]
|
||||
video_positions = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 4, 4, 4, 4, 5, 6, 7],
|
||||
[2, 3, 4, 4, 5, 5, 4, 5, 6, 7],
|
||||
[2, 3, 4, 5, 4, 5, 4, 5, 6, 7],
|
||||
]
|
||||
)
|
||||
for index in range(num_videos):
|
||||
offset = len(input_tokens) + 1
|
||||
mm_features.extend(
|
||||
[
|
||||
make_video_feature(offset, (1, 2, 2), use_audio_in_video=True),
|
||||
make_audio_feature(offset, audio_feature_length=16),
|
||||
]
|
||||
)
|
||||
input_tokens.extend(video_tokens)
|
||||
expected.append(video_positions + index * 6)
|
||||
|
||||
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
input_tokens.extend([TEXT_TOKEN_ID] * trailing_tokens)
|
||||
next_position = 2 + num_videos * 6
|
||||
expected.append((torch.arange(trailing_tokens) + next_position).expand(3, -1))
|
||||
positions, delta = model.get_mrope_input_positions(input_tokens, mm_features)
|
||||
|
||||
assert positions.shape == (3, len(input_tokens))
|
||||
torch.testing.assert_close(positions, torch.cat(expected, dim=1))
|
||||
assert delta == -4 * num_videos
|
||||
|
||||
|
||||
def test_mrope_positions_image_then_text():
|
||||
|
||||
@@ -2311,8 +2311,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
st = offset + video_len + 1
|
||||
else:
|
||||
audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
|
||||
audio_bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(audio_bos_pos)
|
||||
st_idx += 1
|
||||
|
||||
pos_ids, _ = self._compute_interleaved_positions(st_idx, data)
|
||||
llm_pos_ids_list.append(pos_ids)
|
||||
@@ -2320,7 +2321,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
|
||||
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
llm_pos_ids_list.append(eos_pos)
|
||||
llm_pos_ids_list.append(eos_pos + 1)
|
||||
|
||||
video_len = grid_t * grid_h * grid_w
|
||||
audio_len = self._compute_audio_token_count(
|
||||
|
||||
Reference in New Issue
Block a user