[Core] Optimize Qwen2.5-Omni Token2Wav buffer loading (#7061)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

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 = {

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