diff --git a/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py index 1205dedb6..0809b8056 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py +++ b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py @@ -1,8 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav import ( + Qwen2_5OmniToken2WavForConditionalGenerationVLLM, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class _BufferModel(nn.Module): + def __init__(self): + super().__init__() + self.state_dict_calls = 0 + self.register_buffer("persistent", torch.ones(1)) + self.register_buffer("temporary", torch.ones(1), persistent=False) + self.child = nn.Module() + self.child.register_buffer("persistent", torch.ones(1)) + + def state_dict(self, *args, **kwargs): + self.state_dict_calls += 1 + return super().state_dict(*args, **kwargs) + + +def test_token2wav_buffer_discovery_avoids_state_dict(): + model = _BufferModel() + + buffers = Qwen2_5OmniToken2WavForConditionalGenerationVLLM.find_all_registers(model) + + assert set(buffers) == {"persistent", "child.persistent"} + assert model.state_dict_calls == 0 + + class TestQwen2_5OmniWeightLoading: def test_qkv_weight_mapping(self): mock_quant_description = { diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index 49d3ac114..dacb47b88 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + ############################ # Start Token2Wav # ############################ @@ -1509,19 +1512,19 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): return loaded def find_all_registers(self): - """ - Find all registered buffers in a PyTorch model. - - Args: - Returns: - dict: Dictionary with buffer names as keys and their properties as values - """ + """Find all persistent registered buffers in the model.""" registers = {} - # Get all named buffers - for name, buf in self.named_buffers(): - if name in self.state_dict(): - registers[name] = {"name": name, "buffer": buf} + # Match state_dict() by excluding non-persistent buffers. + non_persistent_names = set() + for module_name, module in self.named_modules(): + for buffer_name in getattr(module, "_non_persistent_buffers_set", set()): + name = f"{module_name}.{buffer_name}" if module_name else buffer_name + non_persistent_names.add(name) + + for name, buffer in self.named_buffers(): + if name not in non_persistent_names: + registers[name] = {"name": name, "buffer": buffer} return registers # remove buffers from the weights and reload them after loading weights