wehub-resource-sync
647 lines
22 KiB
Python
647 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from transformers import Qwen3Config
|
|
|
|
from vllm.model_executor.models.interfaces import SupportsLoRA, supports_lora
|
|
from vllm.model_executor.models.moss_audio import (
|
|
MOSS_AUDIO_BOS_TOKEN,
|
|
MOSS_AUDIO_BOS_TOKEN_ID,
|
|
MOSS_AUDIO_EOS_TOKEN,
|
|
MOSS_AUDIO_EOS_TOKEN_ID,
|
|
MOSS_AUDIO_PLACEHOLDER,
|
|
MOSS_AUDIO_TOKEN,
|
|
MOSS_AUDIO_TOKEN_ID,
|
|
MossAudioConfig,
|
|
MossAudioDummyInputsBuilder,
|
|
MossAudioEncoder,
|
|
MossAudioEncoderConfig,
|
|
MossAudioModel,
|
|
MossAudioMultiModalProcessor,
|
|
MossAudioProcessingInfo,
|
|
MossAudioProcessor,
|
|
MossQwen3ForCausalLM,
|
|
MossQwen3Model,
|
|
)
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
|
|
from vllm.multimodal.inputs import batched_tensors_equal
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
|
|
class _Tokenizer:
|
|
def encode(self, text, add_special_tokens=False):
|
|
del add_special_tokens
|
|
return [ord(char) for char in text]
|
|
|
|
def decode(self, token_ids, **kwargs):
|
|
del kwargs
|
|
return "".join(chr(token_id) for token_id in token_ids)
|
|
|
|
def batch_decode(self, batch_token_ids, **kwargs):
|
|
return [self.decode(token_ids, **kwargs) for token_ids in batch_token_ids]
|
|
|
|
|
|
class _MMConfig:
|
|
enable_mm_embeds = False
|
|
mm_processor_cache_gb = 1
|
|
|
|
def merge_mm_processor_kwargs(self, kwargs):
|
|
return dict(kwargs)
|
|
|
|
def get_limit_per_prompt(self, modality):
|
|
del modality
|
|
return 3
|
|
|
|
|
|
class _ModelConfig:
|
|
def __init__(self):
|
|
self.model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
|
|
self.revision = None
|
|
self.max_model_len = 4096
|
|
self.encoder_config = {}
|
|
self.dtype = torch.float32
|
|
self.hf_config = MossAudioConfig(language_config=Qwen3Config())
|
|
self.multimodal_config = _MMConfig()
|
|
|
|
def get_multimodal_config(self):
|
|
return self.multimodal_config
|
|
|
|
def get_inputs_embeds_size(self):
|
|
return None
|
|
|
|
|
|
class _ProcessingContext:
|
|
def __init__(self):
|
|
self.model_config = _ModelConfig()
|
|
self.tokenizer = _Tokenizer()
|
|
|
|
def get_tokenizer(self):
|
|
return self.tokenizer
|
|
|
|
def get_hf_config(self):
|
|
return self.model_config.hf_config
|
|
|
|
def get_mm_config(self):
|
|
return self.model_config.get_multimodal_config()
|
|
|
|
def get_merged_mm_kwargs(self, kwargs):
|
|
return self.get_mm_config().merge_mm_processor_kwargs(kwargs)
|
|
|
|
def call_hf_processor(self, hf_processor, data, kwargs):
|
|
merged_kwargs = self.get_merged_mm_kwargs(kwargs)
|
|
merged_kwargs.setdefault("return_tensors", "pt")
|
|
return hf_processor(**data, **merged_kwargs)
|
|
|
|
|
|
class _TestMossAudioProcessingInfo(MossAudioProcessingInfo):
|
|
def _get_processor_config_defaults(self):
|
|
return {}
|
|
|
|
|
|
def _vllm_config(tensor_parallel_size=1, pipeline_parallel_size=1, hf_config=None):
|
|
if hf_config is None:
|
|
hf_config = MossAudioConfig(language_config=Qwen3Config())
|
|
return SimpleNamespace(
|
|
model_config=SimpleNamespace(
|
|
hf_config=hf_config,
|
|
multimodal_config=None,
|
|
),
|
|
quant_config=None,
|
|
parallel_config=SimpleNamespace(
|
|
tensor_parallel_size=tensor_parallel_size,
|
|
pipeline_parallel_size=pipeline_parallel_size,
|
|
),
|
|
)
|
|
|
|
|
|
class _FakeAudioEncoder:
|
|
dtype = torch.float32
|
|
|
|
def __init__(self, deepstack_layers=0):
|
|
self.deepstack_layers = deepstack_layers
|
|
self.output_deepstack_hidden_states = None
|
|
self.input_shape = None
|
|
self.feature_lens = None
|
|
|
|
def __call__(self, audio_data, *, feature_lens, output_deepstack_hidden_states):
|
|
self.input_shape = tuple(audio_data.shape)
|
|
self.feature_lens = feature_lens.detach().cpu().clone()
|
|
self.output_deepstack_hidden_states = output_deepstack_hidden_states
|
|
lengths = MossAudioEncoder._compute_downsampled_length(feature_lens)
|
|
hidden_states = torch.ones(1, int(lengths.sum().item()), 8)
|
|
if not output_deepstack_hidden_states:
|
|
return hidden_states, None
|
|
return hidden_states, [
|
|
hidden_states * scale for scale in range(2, 2 + self.deepstack_layers)
|
|
]
|
|
|
|
|
|
def _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=1, tp_rank=0):
|
|
import vllm.model_executor.layers.linear as linear_layers
|
|
import vllm.model_executor.models.moss_audio as moss_audio_module
|
|
import vllm.model_executor.parameter as parameter_module
|
|
|
|
for module in (moss_audio_module, linear_layers, parameter_module):
|
|
monkeypatch.setattr(
|
|
module, "get_tensor_model_parallel_world_size", lambda: tp_size
|
|
)
|
|
monkeypatch.setattr(
|
|
linear_layers, "get_tensor_model_parallel_rank", lambda: tp_rank
|
|
)
|
|
monkeypatch.setattr(
|
|
parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank
|
|
)
|
|
monkeypatch.setattr(
|
|
linear_layers, "tensor_model_parallel_all_reduce", lambda tensor: tensor
|
|
)
|
|
|
|
|
|
def _build_moss_audio_processor(cache=None):
|
|
ctx = _ProcessingContext()
|
|
info = _TestMossAudioProcessingInfo(ctx)
|
|
return (
|
|
MossAudioMultiModalProcessor(
|
|
info,
|
|
MossAudioDummyInputsBuilder(info),
|
|
cache=cache,
|
|
),
|
|
ctx,
|
|
)
|
|
|
|
|
|
def _assert_mm_inputs_equal(left, right):
|
|
assert left["prompt_token_ids"] == right["prompt_token_ids"]
|
|
assert left["mm_hashes"] == right["mm_hashes"]
|
|
|
|
left_placeholder = left["mm_placeholders"]["audio"][0]
|
|
right_placeholder = right["mm_placeholders"]["audio"][0]
|
|
assert left_placeholder.offset == right_placeholder.offset
|
|
assert left_placeholder.length == right_placeholder.length
|
|
assert left_placeholder.is_embed.tolist() == right_placeholder.is_embed.tolist()
|
|
|
|
assert batched_tensors_equal(
|
|
left["mm_kwargs"].get_data(),
|
|
right["mm_kwargs"].get_data(),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("prompt", "prefix"),
|
|
[
|
|
(
|
|
f"before {MOSS_AUDIO_PLACEHOLDER} after",
|
|
[*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
|
|
),
|
|
(
|
|
f"before {MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}"
|
|
f"{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN} after",
|
|
[*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
|
|
),
|
|
("Describe this audio.", [MOSS_AUDIO_BOS_TOKEN_ID]),
|
|
],
|
|
)
|
|
def test_moss_audio_processor_expands_audio_placeholders(prompt, prefix):
|
|
raw_mel_len = 17
|
|
processed = MossAudioProcessor(_Tokenizer())(
|
|
text=prompt, audio=[torch.zeros(160 * raw_mel_len)]
|
|
)
|
|
input_ids = processed["input_ids"][0].tolist()
|
|
|
|
assert input_ids[: len(prefix)] == prefix
|
|
assert input_ids.count(MOSS_AUDIO_BOS_TOKEN_ID) == 1
|
|
assert input_ids.count(MOSS_AUDIO_EOS_TOKEN_ID) == 1
|
|
assert input_ids.count(MOSS_AUDIO_TOKEN_ID) == (
|
|
MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
|
|
)
|
|
assert processed["audio_data"].shape == (1, 128, raw_mel_len)
|
|
assert processed["audio_data_seqlens"].tolist() == [raw_mel_len]
|
|
|
|
|
|
def test_moss_audio_processor_preserves_placeholder_without_audio():
|
|
processed = MossAudioProcessor(_Tokenizer())(
|
|
text=f"before {MOSS_AUDIO_PLACEHOLDER} after"
|
|
)
|
|
|
|
assert processed["input_ids"][0].tolist() == [
|
|
*[ord(char) for char in "before "],
|
|
MOSS_AUDIO_BOS_TOKEN_ID,
|
|
MOSS_AUDIO_TOKEN_ID,
|
|
MOSS_AUDIO_EOS_TOKEN_ID,
|
|
*[ord(char) for char in " after"],
|
|
]
|
|
assert "audio_data" not in processed
|
|
assert "audio_data_seqlens" not in processed
|
|
|
|
|
|
def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
|
|
raw_mel_len = 17
|
|
audio = np.zeros(160 * raw_mel_len, dtype=np.float32)
|
|
prompt = f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."
|
|
|
|
baseline_processor, ctx = _build_moss_audio_processor()
|
|
mm_items = baseline_processor.info.parse_mm_data({"audio": [audio]})
|
|
token_prompt = ctx.get_tokenizer().encode(prompt, add_special_tokens=False)
|
|
|
|
baseline_text = baseline_processor(
|
|
prompt,
|
|
mm_items=mm_items,
|
|
hf_processor_mm_kwargs={},
|
|
)
|
|
baseline_token = baseline_processor(
|
|
token_prompt,
|
|
mm_items=mm_items,
|
|
hf_processor_mm_kwargs={},
|
|
)
|
|
|
|
cache = MultiModalProcessorOnlyCache(ctx.model_config)
|
|
cached_processor, _ = _build_moss_audio_processor(cache=cache)
|
|
cached_text_miss = cached_processor(
|
|
prompt,
|
|
mm_items=mm_items,
|
|
hf_processor_mm_kwargs={},
|
|
)
|
|
cached_text_hit = cached_processor(
|
|
prompt,
|
|
mm_items=mm_items,
|
|
hf_processor_mm_kwargs={},
|
|
)
|
|
cached_token_hit = cached_processor(
|
|
token_prompt,
|
|
mm_items=mm_items,
|
|
hf_processor_mm_kwargs={},
|
|
)
|
|
|
|
expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
|
|
prompt_token_ids = baseline_text["prompt_token_ids"]
|
|
assert prompt_token_ids.count(MOSS_AUDIO_TOKEN_ID) == expected_audio_tokens
|
|
assert baseline_text["mm_placeholders"]["audio"][0].length == (
|
|
expected_audio_tokens + 2
|
|
)
|
|
|
|
_assert_mm_inputs_equal(baseline_text, baseline_token)
|
|
_assert_mm_inputs_equal(baseline_text, cached_text_miss)
|
|
_assert_mm_inputs_equal(baseline_text, cached_text_hit)
|
|
_assert_mm_inputs_equal(baseline_text, cached_token_hit)
|
|
|
|
|
|
def test_moss_audio_supports_language_model_lora_only():
|
|
assert supports_lora(MossAudioModel)
|
|
|
|
model = object.__new__(MossAudioModel)
|
|
assert isinstance(model, SupportsLoRA)
|
|
|
|
mapping = model.get_mm_mapping()
|
|
assert mapping.language_model == ["language_model."]
|
|
assert mapping.tower_model == []
|
|
assert mapping.connector == []
|
|
|
|
|
|
def test_moss_audio_error_paths():
|
|
model = object.__new__(MossAudioModel)
|
|
with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
|
|
model._cache_deepstack_input_embeds(
|
|
inputs_embeds=torch.zeros(4, 8),
|
|
deepstack_embeddings=((torch.ones(1, 8),),),
|
|
is_multimodal=torch.tensor([False, True, True, False]),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="too short"):
|
|
MossAudioProcessor(_Tokenizer())(
|
|
text=MOSS_AUDIO_PLACEHOLDER, audio=[torch.empty(0)]
|
|
)
|
|
with pytest.raises(ValueError, match="too short"):
|
|
model._parse_and_validate_audio_input(
|
|
audio_data=torch.zeros(1, 128, 1),
|
|
audio_data_seqlens=torch.tensor([0], dtype=torch.long),
|
|
)
|
|
|
|
|
|
def test_moss_audio_validates_tp_config():
|
|
vllm_config = _vllm_config(tensor_parallel_size=2)
|
|
vllm_config.model_config.hf_config.adapter_hidden_size = 7
|
|
|
|
with pytest.raises(ValueError, match="adapter_hidden_size"):
|
|
MossAudioModel(vllm_config=vllm_config)
|
|
|
|
vllm_config = _vllm_config(tensor_parallel_size=2)
|
|
vllm_config.model_config.hf_config.audio_config.d_model = 6
|
|
vllm_config.model_config.hf_config.audio_config.encoder_attention_heads = 3
|
|
with pytest.raises(ValueError, match="encoder_attention_heads"):
|
|
MossAudioModel(vllm_config=vllm_config)
|
|
|
|
|
|
def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
|
|
model = object.__new__(MossAudioModel)
|
|
|
|
with pytest.raises(ValueError, match="audio_data batch size"):
|
|
model._parse_and_validate_audio_input(
|
|
audio_data=[torch.zeros(128, 8), torch.zeros(128, 11)],
|
|
audio_data_seqlens=torch.tensor([8], dtype=torch.long),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
|
|
def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
|
|
model = object.__new__(MossAudioModel)
|
|
model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
|
|
model.audio_adapter = lambda hidden_states: hidden_states * 5
|
|
model.deepstack_audio_merger_list = [
|
|
lambda hidden_states, scale=scale: hidden_states * scale
|
|
for scale in deepstack_scales
|
|
]
|
|
model.deepstack_input_embeds = None
|
|
|
|
embeddings = model.embed_multimodal(
|
|
audio_data=torch.zeros(2, 128, 9),
|
|
audio_data_seqlens=torch.tensor([8, 9], dtype=torch.long),
|
|
)
|
|
|
|
assert model.audio_encoder.output_deepstack_hidden_states is bool(deepstack_scales)
|
|
assert [embeds.shape for embeds in embeddings] == [
|
|
torch.Size([1, 8 * (1 + len(deepstack_scales))]),
|
|
torch.Size([2, 8 * (1 + len(deepstack_scales))]),
|
|
]
|
|
if not deepstack_scales:
|
|
assert model.deepstack_input_embeds is None
|
|
return
|
|
|
|
main_embeddings, deepstack_embeddings = model._split_multimodal_embeddings(
|
|
embeddings, hidden_size=8
|
|
)
|
|
assert [embeds.shape for embeds in main_embeddings] == [
|
|
torch.Size([1, 8]),
|
|
torch.Size([2, 8]),
|
|
]
|
|
assert [[e.shape for e in layer] for layer in deepstack_embeddings] == [
|
|
[torch.Size([1, 8]), torch.Size([2, 8])] for _ in deepstack_scales
|
|
]
|
|
assert torch.equal(main_embeddings[0], torch.full((1, 8), 5.0))
|
|
for idx, scale in enumerate(deepstack_scales):
|
|
assert torch.equal(
|
|
deepstack_embeddings[idx][0],
|
|
torch.full((1, 8), float((idx + 2) * scale)),
|
|
)
|
|
|
|
|
|
def test_moss_audio_embed_input_ids_caches_packed_deepstack():
|
|
class _FakeLanguageModel:
|
|
def embed_input_ids(self, input_ids):
|
|
return torch.zeros(input_ids.shape[0], 8)
|
|
|
|
model = object.__new__(MossAudioModel)
|
|
model.language_model = _FakeLanguageModel()
|
|
model.deepstack_audio_merger_list = [object(), object()]
|
|
model.deepstack_input_embeds = None
|
|
multimodal_embeddings = (
|
|
torch.cat([torch.full((1, 8), x) for x in (5.0, 14.0, 33.0)], dim=-1),
|
|
torch.cat([torch.full((2, 8), x) for x in (7.0, 22.0, 44.0)], dim=-1),
|
|
)
|
|
is_multimodal = torch.tensor([False, True, True, True, False])
|
|
|
|
inputs_embeds = model.embed_input_ids(
|
|
input_ids=torch.arange(5),
|
|
multimodal_embeddings=multimodal_embeddings,
|
|
is_multimodal=is_multimodal,
|
|
)
|
|
|
|
assert torch.equal(inputs_embeds[1], torch.full((8,), 5.0))
|
|
assert torch.equal(inputs_embeds[2], torch.full((8,), 7.0))
|
|
assert torch.equal(inputs_embeds[3], torch.full((8,), 7.0))
|
|
assert model.deepstack_input_embeds is not None
|
|
tensors = model.deepstack_input_embeds.tensors
|
|
assert set(tensors) == {"deepstack_input_embeds_0", "deepstack_input_embeds_1"}
|
|
for tensor in tensors.values():
|
|
assert tensor[is_multimodal].abs().sum() > 0
|
|
assert torch.equal(tensor[~is_multimodal], torch.zeros(2, 8))
|
|
|
|
|
|
def _patch_pp_group(monkeypatch, *, first=True, last=True):
|
|
import vllm.model_executor.models.moss_audio as moss_audio_module
|
|
|
|
monkeypatch.setattr(
|
|
moss_audio_module,
|
|
"get_pp_group",
|
|
lambda: SimpleNamespace(is_first_rank=first, is_last_rank=last),
|
|
)
|
|
|
|
|
|
def test_moss_audio_pp_forward_routes_deepstack(monkeypatch):
|
|
for first in (True, False):
|
|
calls: list[dict[str, object]] = []
|
|
|
|
def fake_lm(*args, _calls=calls, **kwargs):
|
|
del args
|
|
_calls.append(kwargs)
|
|
return torch.ones(1, 1)
|
|
|
|
_patch_pp_group(monkeypatch, first=first)
|
|
model = object.__new__(MossAudioModel)
|
|
torch.nn.Module.__init__(model)
|
|
model.language_model = fake_lm
|
|
cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
|
|
inter = IntermediateTensors(
|
|
{
|
|
"hidden_states": torch.ones(3, 8),
|
|
"residual": torch.zeros(3, 8),
|
|
"deepstack_input_embeds_0": torch.full((3, 8), 5.0),
|
|
}
|
|
)
|
|
inputs_embeds = torch.full((3, 8), 9.0)
|
|
model.deepstack_input_embeds = cached
|
|
|
|
model.forward(
|
|
input_ids=None,
|
|
positions=torch.arange(3),
|
|
intermediate_tensors=None if first else inter,
|
|
inputs_embeds=inputs_embeds if first else None,
|
|
)
|
|
|
|
kwargs = calls[0]
|
|
assert kwargs["inputs_embeds"] is (inputs_embeds if first else None)
|
|
assert kwargs["deepstack_input_embeds"] is (cached if first else inter)
|
|
assert model.deepstack_input_embeds is None
|
|
|
|
calls = []
|
|
|
|
def fake_lm_non_first_rank(*args, **kwargs):
|
|
del args
|
|
calls.append(kwargs)
|
|
return torch.ones(1, 1)
|
|
|
|
_patch_pp_group(monkeypatch, first=False)
|
|
model = object.__new__(MossAudioModel)
|
|
torch.nn.Module.__init__(model)
|
|
model.language_model = fake_lm_non_first_rank
|
|
model.deepstack_input_embeds = IntermediateTensors({})
|
|
inter = IntermediateTensors(
|
|
{
|
|
"hidden_states": torch.ones(3, 8),
|
|
"residual": torch.zeros(3, 8),
|
|
}
|
|
)
|
|
|
|
model.forward(
|
|
input_ids=None,
|
|
positions=torch.arange(3),
|
|
intermediate_tensors=inter,
|
|
inputs_embeds=torch.ones(3, 8),
|
|
)
|
|
assert calls[0]["inputs_embeds"] is None
|
|
assert calls[0]["deepstack_input_embeds"] is inter
|
|
|
|
|
|
def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
|
|
class AddOne(torch.nn.Module):
|
|
def forward(self, positions, hidden_states, residual):
|
|
del positions, residual
|
|
return hidden_states + 1, torch.zeros_like(hidden_states)
|
|
|
|
def make_model(num_layers, deepstack_layers=None):
|
|
model = object.__new__(MossQwen3Model)
|
|
torch.nn.Module.__init__(model)
|
|
model.start_layer, model.end_layer = 0, num_layers
|
|
model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
|
|
model.norm = lambda hidden_states, residual: (hidden_states, residual)
|
|
model._maybe_add_hidden_state = lambda aux, *args: aux
|
|
model.deepstack_inject_layer_indices = (
|
|
range(0) if deepstack_layers is None else deepstack_layers
|
|
)
|
|
return model
|
|
|
|
_patch_pp_group(monkeypatch, first=True, last=True)
|
|
output = make_model(3).forward(
|
|
input_ids=None,
|
|
positions=torch.arange(2),
|
|
inputs_embeds=torch.zeros(2, 4),
|
|
deepstack_input_embeds=IntermediateTensors(
|
|
{
|
|
"deepstack_input_embeds_2": torch.full((2, 4), 5.0),
|
|
}
|
|
),
|
|
)
|
|
assert torch.equal(output, torch.full((2, 4), 8.0))
|
|
|
|
_patch_pp_group(monkeypatch, first=True, last=False)
|
|
deepstack = IntermediateTensors(
|
|
{
|
|
"deepstack_input_embeds_0": torch.full((2, 4), 7.0),
|
|
"deepstack_input_embeds_3": torch.full((2, 4), 11.0),
|
|
}
|
|
)
|
|
output = make_model(2, range(4)).forward(
|
|
input_ids=None,
|
|
positions=torch.arange(2),
|
|
inputs_embeds=torch.zeros(2, 4),
|
|
deepstack_input_embeds=deepstack,
|
|
)
|
|
assert isinstance(output, IntermediateTensors)
|
|
assert set(output.tensors) == {
|
|
"hidden_states",
|
|
"residual",
|
|
"deepstack_input_embeds_2",
|
|
"deepstack_input_embeds_3",
|
|
}
|
|
assert torch.equal(output["hidden_states"], torch.full((2, 4), 9.0))
|
|
assert torch.equal(output["deepstack_input_embeds_2"], torch.zeros(2, 4))
|
|
assert output["deepstack_input_embeds_3"] is deepstack["deepstack_input_embeds_3"]
|
|
|
|
inner_model = make_model(0, range(2))
|
|
inner_model.make_empty_intermediate_tensors = lambda batch, dtype, device: (
|
|
IntermediateTensors(
|
|
{
|
|
"hidden_states": torch.zeros(batch, 4, dtype=dtype, device=device),
|
|
"residual": torch.zeros(batch, 4, dtype=dtype, device=device),
|
|
}
|
|
)
|
|
)
|
|
language_model = object.__new__(MossQwen3ForCausalLM)
|
|
torch.nn.Module.__init__(language_model)
|
|
language_model.model = inner_model
|
|
language_model.config = SimpleNamespace(hidden_size=4)
|
|
language_model.deepstack_inject_layer_indices = range(2)
|
|
|
|
tensors = MossQwen3ForCausalLM.make_empty_intermediate_tensors(
|
|
language_model,
|
|
batch_size=3,
|
|
dtype=torch.float16,
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
assert set(tensors.tensors) == {
|
|
"hidden_states",
|
|
"residual",
|
|
"deepstack_input_embeds_0",
|
|
"deepstack_input_embeds_1",
|
|
}
|
|
assert tensors["deepstack_input_embeds_0"].shape == (3, 4)
|
|
assert tensors["deepstack_input_embeds_0"].dtype == torch.float16
|
|
|
|
_patch_pp_group(monkeypatch, first=True, last=False)
|
|
forward_tensors = inner_model.forward(
|
|
input_ids=None,
|
|
positions=torch.arange(3),
|
|
inputs_embeds=torch.ones(3, 4, dtype=torch.float16),
|
|
deepstack_input_embeds=None,
|
|
)
|
|
assert isinstance(forward_tensors, IntermediateTensors)
|
|
assert set(forward_tensors.tensors) == set(tensors.tensors)
|
|
|
|
|
|
def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.config.device import DeviceConfig
|
|
|
|
_patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=2)
|
|
config = MossAudioEncoderConfig(
|
|
d_model=8,
|
|
output_dim=8,
|
|
num_mel_bins=8,
|
|
encoder_layers=1,
|
|
encoder_attention_heads=2,
|
|
encoder_ffn_dim=16,
|
|
downsample_hidden_size=2,
|
|
deepstack_encoder_layer_indexes=[],
|
|
)
|
|
with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
|
|
encoder = MossAudioEncoder(config)
|
|
|
|
attention = encoder.layers[0].self_attn
|
|
assert all(hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj"))
|
|
assert hasattr(attention, "out_proj")
|
|
assert not hasattr(attention, "qkv")
|
|
assert attention.k_proj.bias is None
|
|
|
|
weight_names = [
|
|
"layers.0.self_attn.q_proj.weight",
|
|
"layers.0.self_attn.q_proj.bias",
|
|
"layers.0.self_attn.k_proj.weight",
|
|
"layers.0.self_attn.v_proj.weight",
|
|
"layers.0.self_attn.v_proj.bias",
|
|
"layers.0.self_attn.out_proj.weight",
|
|
"layers.0.self_attn.out_proj.bias",
|
|
"conv1.weight",
|
|
"conv1.bias",
|
|
]
|
|
params = dict(encoder.named_parameters(remove_duplicate=False))
|
|
assert "layers.0.self_attn.k_proj.bias" not in params
|
|
weights = {
|
|
name: torch.full_like(params[name], fill_value=float(i + 1))
|
|
for i, name in enumerate(weight_names)
|
|
}
|
|
|
|
loaded = AutoWeightsLoader(encoder).load_weights(weights.items())
|
|
|
|
assert "load_weights" not in MossAudioEncoder.__dict__
|
|
assert loaded == set(weight_names)
|
|
assert not any(".qkv." in name for name in loaded)
|
|
assert torch.equal(
|
|
params["layers.0.self_attn.q_proj.weight"],
|
|
weights["layers.0.self_attn.q_proj.weight"],
|
|
)
|