diff --git a/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_ar_kv_layers.py b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_ar_kv_layers.py new file mode 100644 index 000000000..cc9275894 --- /dev/null +++ b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_ar_kv_layers.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AR stage must register attention layers for its own language model only. + +The Qwen-VL parent ``__init__`` builds a stock language model before MammothModa2 +replaces it. Its attention layers stay in ``static_forward_context`` unless they are +removed, and the KV cache is then sized and allocated for both models. The parent +``__init__`` and the replacement are stubbed here so only the registry is exercised. +""" + +import gc +import weakref +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.model_executor.models.utils import StageMissingLayer + +import vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2 as mm2 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_LAYERS = 3 +VISION_LAYER = "ar.visual.blocks.0.attn" +AR_CLASSES = [ + (mm2.MammothModa2ARForConditionalGeneration, mm2.Qwen2_5_VLForConditionalGeneration), + (mm2.MammothModa2Qwen3ARForConditionalGeneration, mm2.Qwen3VLForConditionalGeneration), +] +AR_IDS = ["preview_qwen2_5_vl", "dev_qwen3_vl"] + + +class _FakeAttention(nn.Module): + def __init__(self, registry: dict, name: str): + super().__init__() + registry[name] = self + + +class _FakeLM(nn.Module): + def __init__(self, registry: dict, prefix: str): + super().__init__() + self.layers = nn.ModuleList( + _FakeAttention(registry, f"{prefix}.layers.{i}.self_attn.attn") for i in range(NUM_LAYERS) + ) + self.make_empty_intermediate_tensors = lambda *args, **kwargs: None + + +class _FakeCompiledLM(_FakeLM): + cleanups = 0 + + def cleanup(self): + type(self).cleanups += 1 + + +def _build( + monkeypatch, + ar_class, + parent, + *, + encoder_only=False, + parent_lm=_FakeLM, + parent_registers=True, + pp_rank=(1, True), + stray_parent_layer=False, +): + registry: dict = {} + built = {} + + def parent_init(self, *, vllm_config, prefix=""): + nn.Module.__init__(self) + self.visual = _FakeAttention(registry, VISION_LAYER) + stock = nn.Module() + stock.model = parent_lm(registry if parent_registers else {}, f"{prefix}.language_model.model") + built["parent_lm"] = weakref.ref(stock) + if stray_parent_layer: + # A parent-prefixed layer that is not part of the module being dropped. + registry[f"{prefix}.language_model.model.layers.9.self_attn.attn"] = nn.Module() + self.language_model = StageMissingLayer("language_model", stock) if encoder_only else stock + + # The model module only imports get_current_vllm_config with this fix; keep the + # reverted source testable so the registry assertion is what fails there. + monkeypatch.setattr( + mm2, + "get_current_vllm_config", + lambda: SimpleNamespace(compilation_config=SimpleNamespace(static_forward_context=registry)), + raising=False, + ) + monkeypatch.setattr(parent, "__init__", parent_init) + monkeypatch.setattr(mm2, "init_vllm_registered_model", lambda *, prefix, **kwargs: _FakeLM(registry, prefix)) + monkeypatch.setattr(mm2, "MammothModa2Qwen3ForCausalLM", lambda *, vllm_config, prefix: _FakeLM(registry, prefix)) + monkeypatch.setattr(mm2, "TorchCompileWithNoGuardsWrapper", _FakeCompiledLM, raising=False) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + world_size, is_first_rank = pp_rank + monkeypatch.setattr( + mm2, "get_pp_group", lambda: SimpleNamespace(world_size=world_size, is_first_rank=is_first_rank) + ) + + hf_config = SimpleNamespace(llm_config=SimpleNamespace(text_config=SimpleNamespace())) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=hf_config, architectures=["MammothModa2Model"]) + ) + vllm_config.with_hf_config = lambda hf, architectures=None: vllm_config + model = ar_class(vllm_config=vllm_config, prefix="ar") + return model, registry, built + + +@pytest.mark.parametrize("encoder_only", [False, True], ids=["default", "mm_encoder_only"]) +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_ar_stage_registers_only_the_replacement_language_model(monkeypatch, ar_class, parent, encoder_only): + model, registry, built = _build(monkeypatch, ar_class, parent, encoder_only=encoder_only) + + assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER} + gc.collect() + assert built["parent_lm"]() is None + + +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_compiled_parent_language_model_releases_its_hook(monkeypatch, ar_class, parent): + _FakeCompiledLM.cleanups = 0 + _build(monkeypatch, ar_class, parent, parent_lm=_FakeCompiledLM) + assert _FakeCompiledLM.cleanups == 1 + + +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_unexpected_parent_registry_fails_loudly(monkeypatch, ar_class, parent): + with pytest.raises(RuntimeError, match="Unexpected attention registry"): + _build(monkeypatch, ar_class, parent, parent_registers=False) + + +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_later_pipeline_rank_without_parent_layers_starts(monkeypatch, ar_class, parent): + _, registry, _ = _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False)) + assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER} + + +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_first_pipeline_rank_without_parent_layers_fails(monkeypatch, ar_class, parent): + with pytest.raises(RuntimeError, match="Unexpected attention registry"): + _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, True)) + + +@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS) +def test_later_pipeline_rank_still_rejects_a_leftover_parent_layer(monkeypatch, ar_class, parent): + with pytest.raises(RuntimeError, match="Unexpected attention registry"): + _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False), stray_parent_layer=True) diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py index 68c4df022..885d4b0df 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py @@ -1,5 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +import gc from collections.abc import Callable, Iterable, Mapping from itertools import islice from typing import Any @@ -8,8 +12,10 @@ import torch from torch import nn from transformers import Qwen2Config, Qwen3VLProcessor from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor -from vllm.config import CacheConfig, VllmConfig +from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import get_pp_group +from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -40,6 +46,7 @@ from vllm.model_executor.models.qwen3_vl import ( from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + StageMissingLayer, WeightsMapper, init_vllm_registered_model, is_pp_missing_parameter, @@ -57,6 +64,8 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, is_interleaved from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config +logger = init_logger(__name__) + def _runtime_meta(runtime_info: dict[str, Any]) -> dict[str, Any]: """Accept both legacy flat metadata and the canonical nested payload.""" @@ -623,6 +632,42 @@ class MammothModa2Qwen3ForCausalLM(MammothModa2Qwen2ForCausalLM): ) +def _drop_parent_language_model(model: nn.Module, prefix: str) -> None: + """Unregister and release the language model built by the Qwen-VL parent ``__init__``. + + The parent builds a stock Qwen LM whose attention layers register in the static + forward context as ``{prefix}.language_model.model.layers.*``. MammothModa2 replaces + it with its own LM (``{prefix}.language_model.layers.*``); without this, both sets of + layers are counted when the KV cache is sized and each gets its own KV blocks. + """ + old = model.language_model + if isinstance(old, StageMissingLayer): # mm_encoder_only + old = old.__dict__["module"] + ctx = get_current_vllm_config().compilation_config.static_forward_context + owned = set(old.modules()) + dropped = [name for name, layer in ctx.items() if layer in owned] + for name in dropped: + del ctx[name] + stale = [name for name in ctx if name.startswith(maybe_prefix(prefix, "language_model."))] + # A later pipeline rank may hold none of the parent's layers, so only the first + # rank has to find some; no rank may leave any behind. + pp_group = get_pp_group() + must_drop = pp_group.world_size == 1 or pp_group.is_first_rank + if (must_drop and not dropped) or stale: + raise RuntimeError( + f"Unexpected attention registry for the parent language model (dropped {len(dropped)}, left {stale}); " + "check the Qwen-VL __init__." + ) + # With torch.compile, the compile wrapper's bytecode hook would keep the module alive. + for module in owned: + if isinstance(module, TorchCompileWithNoGuardsWrapper): + module.cleanup() + logger.debug("Dropped %d attention layers of the parent language model.", len(dropped)) + del model.language_model, old, owned + gc.collect() + torch.accelerator.empty_cache() + + @MULTIMODAL_REGISTRY.register_processor( MammothModa2ARMultiModalProcessor, info=MammothModa2ARProcessingInfo, @@ -656,6 +701,7 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration) ar_vllm_config = vllm_config.with_hf_config(ar_hf_config, architectures=vllm_config.model_config.architectures) # Initialize multi-modal components like the vision tower first. super().__init__(vllm_config=ar_vllm_config, prefix=prefix) + _drop_parent_language_model(self, prefix) # Replace with the custom MoE language model. lm_hf_config = getattr( ar_vllm_config.model_config.hf_config, "text_config", ar_vllm_config.model_config.hf_config @@ -824,6 +870,7 @@ class MammothModa2Qwen3ARForConditionalGeneration(Qwen3VLForConditionalGeneratio architectures=vllm_config.model_config.architectures, ) super().__init__(vllm_config=ar_vllm_config, prefix=prefix) + _drop_parent_language_model(self, prefix) text_config = ar_hf_config.text_config self.language_model = MammothModa2Qwen3ForCausalLM( vllm_config=ar_vllm_config.with_hf_config(text_config),