Repository: vllm-project/vllm

Issue:
[Bug]: ModelRegistry.load_model_cls() circular import error on llama-llava
### Your current environment

```text
datta0@naillm:~/kserve/python/huggingfaceserver$ python collect_env.py 
Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.29.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-101-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-PCIE-40GB
Nvidia driver version: 525.60.13
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             8
On-line CPU(s) list:                0-7
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Gold 6354 CPU @ 3.00GHz
CPU family:                         15
Model:                              6
Thread(s) per core:                 1
Core(s) per socket:                 1
Socket(s):                          8
Stepping:                           1
BogoMIPS:                           6000.01
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss syscall nx pdpe1gb rdtscp lm constant_tsc nopl xtopology cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm md_clear arch_capabilities
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          256 KiB (8 instances)
L1i cache:                          256 KiB (8 instances)
L2 cache:                           32 MiB (8 instances)
L3 cache:                           128 MiB (8 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-7
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Unknown: No mitigations
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] flake8==7.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-nccl-cu12==2.20.5
[pip3] torch==2.3.0
[pip3] torchvision==0.16.2
[pip3] triton==2.3.0
[pip3] vllm_nccl_cu12==2.18.1.0.4.0
[conda] Could not collectROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.4.2
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0    CPU Affinity    NUMA Affinity
GPU0     X      0-7             N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks
```


### 🐛 Describe the bug

Steps to reproduce:
1. Install vllm 0.4.2 (works fine on 0.4.0)
2. Run the following code

```python
from vllm.model_executor.models import ModelRegistry
architecture = 'MistralForCausalLM' # can also use LlamaForCausalLM
model_cls = ModelRegistry.load_model_cls(architecture)
```

This throws the error (click to expand)

<details>
<summary>ImportError: cannot import name 'LlamaModel' from partially initialized module 'vllm.model_executor.models.llama' (most likely due to a circular import) (/home/datta0/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/llama.py)"</summary>
```
{
	"name": "ImportError",
	"message": "cannot import name 'LlamaModel' from partially initialized module 'vllm.model_executor.models.llama' (most likely due to a circular import) (/home/datta0/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/llama.py)",
	"stack": "---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], line 3
      1 from vllm.model_executor.models import ModelRegistry
      2 architecture = 'MistralForCausalLM'
----> 3 model_cls = ModelRegistry.load_model_cls(architecture)

File ~/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/__init__.py:98, in ModelRegistry.load_model_cls(model_arch)
     93         logger.warning(
     94             \"Model architecture %s is partially supported by ROCm: %s\",
     95             model_arch, _ROCM_PARTIALLY_SUPPORTED_MODELS[model_arch])
     97 module_name, model_cls_name = _MODELS[model_arch]
---> 98 module = importlib.import_module(
     99     f\"vllm.model_executor.models.{module_name}\")
    100 return getattr(module, model_cls_name, None)

File /usr/lib/python3.10/importlib/__init__.py:126, in import_module(name, package)
    124             break
    125         level += 1
--> 126 return _bootstrap._gcd_import(name[level:], package, level)

File <frozen importlib._bootstrap>:1050, in _gcd_import(name, package, level)

File <frozen importlib._bootstrap>:1027, in _find_and_load(name, import_)

File <frozen importlib._bootstrap>:1006, in _find_and_load_unlocked(name, import_)

File <frozen importlib._bootstrap>:688, in _load_unlocked(spec)

File <frozen importlib._bootstrap_external>:883, in exec_module(self, module)

File <frozen importlib._bootstrap>:241, in _call_with_frames_removed(f, *args, **kwds)

File ~/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/llama.py:46
     43 from vllm.model_executor.layers.sampler import Sampler
     44 from vllm.model_executor.layers.vocab_parallel_embedding import (
     45     DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
---> 46 from vllm.model_executor.model_loader.weight_utils import (
     47     default_weight_loader, kv_cache_scales_loader)
     48 from vllm.model_executor.sampling_metadata import SamplingMetadata
     49 from vllm.sequence import SamplerOutput

File ~/.pyenv/lib/python3.10/site-packages/vllm/model_executor/model_loader/__init__.py:7
      3 from torch import nn
      5 from vllm.config import (DeviceConfig, LoadConfig, LoRAConfig, ModelConfig,
      6                          ParallelConfig, SchedulerConfig, VisionLanguageConfig)
----> 7 from vllm.model_executor.model_loader.loader import (BaseModelLoader,
      8                                                      get_model_loader)
      9 from vllm.model_executor.model_loader.utils import (
     10     get_architecture_class_name, get_model_architecture)
     13 def get_model(
     14         *, model_config: ModelConfig, load_config: LoadConfig,
     15         device_config: DeviceConfig, parallel_config: ParallelConfig,
     16         scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig],
     17         vision_language_config: Optional[VisionLanguageConfig]) -> nn.Module:

File ~/.pyenv/lib/python3.10/site-packages/vllm/model_executor/model_loader/loader.py:28
     22 from vllm.model_executor.model_loader.utils import (get_model_architecture,
     23                                                     set_default_torch_dtype)
     24 from vllm.model_executor.model_loader.weight_utils import (
     25     download_weights_from_hf, filter_files_not_needed_for_inference,
     26     get_quant_config, initialize_dummy_weights, np_cache_weights_iterator,
     27     pt_weights_iterator, safetensors_weights_iterator)
---> 28 from vllm.model_executor.models.llava import LlavaForConditionalGeneration
     30 _VISION_MODEL_CLASSES = [
     31     LlavaForConditionalGeneration,
     32 ]
     34 logger = init_logger(__name__)

File ~/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/llava.py:18
     16 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
     17 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
---> 18 from vllm.model_executor.models.llama import LlamaModel
     19 from vllm.model_executor.sampling_metadata import SamplingMetadata
     20 from vllm.sequence import SamplerOutput

ImportError: cannot import name 'LlamaModel' from partially initialized module 'vllm.model_executor.models.llama' (most likely due to a circular import) (/home/datta0/.pyenv/lib/python3.10/site-packages/vllm/model_executor/models/llama.py)"
}```
</details>


## Cause of the issue: 
[Mistral (LlaMA) executor](https://github.com/vllm-project/vllm/blob/c579b750a083931ad03ecac898aca5ad67c6c59c/vllm/model_executor/models/llama.py#L46) imports model_loader whose [imports called BaseModelLoader](https://github.com/vllm-project/vllm/blob/c579b750a083931ad03ecac898aca5ad67c6c59c/vllm/model_executor/model_loader/__init__.py#L8) which in turn [uses LlavaForConditionalGeneration](https://github.com/vllm-project/vllm/blob/c579b750a083931ad03ecac898aca5ad67c6c59c/vllm/model_executor/model_loader/loader.py#L28) that depends on [llama](https://github.com/vllm-project/vllm/blob/c579b750a083931ad03ecac898aca5ad67c6c59c/vllm/model_executor/models/llava.py#L18) creating a circular dependency.






Discussion:

Resolving patch:
diff --git a/vllm/model_executor/model_loader/loader.py b/vllm/model_executor/model_loader/loader.py
--- a/vllm/model_executor/model_loader/loader.py
+++ b/vllm/model_executor/model_loader/loader.py
@@ -26,11 +26,7 @@
     download_weights_from_hf, filter_files_not_needed_for_inference,
     get_quant_config, initialize_dummy_weights, np_cache_weights_iterator,
     pt_weights_iterator, safetensors_weights_iterator)
-from vllm.model_executor.models.llava import LlavaForConditionalGeneration
-
-_VISION_MODEL_CLASSES = [
-    LlavaForConditionalGeneration,
-]
+from vllm.model_executor.models.vlm_base import VisionLanguageModelBase
 
 logger = init_logger(__name__)
 
@@ -73,7 +69,12 @@ def _get_model_initialization_kwargs(
             "but LoRA is enabled. Support for this model may "
             "be added in the future. If this is important to you, "
             "please open an issue on github.")
-    elif model_class in _VISION_MODEL_CLASSES:
+    elif issubclass(model_class, VisionLanguageModelBase):
+        if vision_language_config is None:
+            raise ValueError("Provide `image_input_type` and other vision "
+                             "related configurations through LLM entrypoint "
+                             "or engine arguments.")
+
         extra_kwargs["vision_language_config"] = vision_language_config
     return extra_kwargs
 
diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py
--- a/vllm/model_executor/models/llava.py
+++ b/vllm/model_executor/models/llava.py
@@ -19,6 +19,8 @@
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import SamplerOutput
 
+from .vlm_base import VisionLanguageModelBase
+
 _KEYS_TO_MODIFY_MAPPING = {
     "language_model.lm_head": "lm_head",
     "language_model.model": "language_model",
@@ -40,7 +42,7 @@ def __init__(self, vision_hidden_size: int, text_hidden_size: int,
                                   text_hidden_size,
                                   bias=True)
 
-    def forward(self, image_features):
+    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
         hidden_states = self.linear_1(image_features)
         hidden_states = self.act(hidden_states)
         hidden_states = self.linear_2(hidden_states)
@@ -50,29 +52,31 @@ def forward(self, image_features):
 def _merge_vision_embeddings(input_ids: torch.Tensor,
                              inputs_embeds: torch.Tensor,
                              vision_embeddings: torch.Tensor,
-                             image_token_id: int):
+                             image_token_id: int) -> torch.Tensor:
     """In place merges in vision_embeddings with inputs_embeds."""
     mask = (input_ids == image_token_id)
-    inputs_embeds[mask] = vision_embeddings.view(-1,
+
+    image_feature_size = vision_embeddings.shape[0] * vision_embeddings.shape[1]
+    if mask.sum() != image_feature_size:
+        raise ValueError(f"image_feature_size should be {image_feature_size}, "
+                         f"but found: {mask.sum()}")
+
+    inputs_embeds[mask] = vision_embeddings.view(image_feature_size,
                                                  vision_embeddings.shape[-1])
 
+    return inputs_embeds
 
-class LlavaForConditionalGeneration(nn.Module):
+
+class LlavaForConditionalGeneration(VisionLanguageModelBase):
 
     def __init__(self,
-                 config: "LlavaConfig",
+                 config: LlavaConfig,
                  vision_language_config: VisionLanguageConfig,
                  cache_config: Optional[CacheConfig] = None,
-                 quant_config: Optional["QuantizationConfig"] = None) -> None:
-        super().__init__()
-        self.config = config
-
-        self.vision_language_config = vision_language_config
+                 quant_config: Optional[QuantizationConfig] = None) -> None:
+        super().__init__(vision_language_config)
 
-        assert self.vision_language_config, (
-            "Provide `image_input_type` and other vision "
-            "related configurations through LLM entrypoint "
-            "or engine arguments.")
+        self.config = config
 
         if self.vision_language_config.image_input_type == (
                 VisionLanguageConfig.ImageInputType.PIXEL_VALUES):
@@ -98,14 +102,12 @@ def __init__(self,
                                                 config.vocab_size, logit_scale)
         self.sampler = Sampler()
 
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        kv_caches: List[torch.Tensor],
-        attn_metadata: AttentionMetadata,
-        image_input: Optional[torch.Tensor] = None
-    ) -> SamplerOutput:  # noqa: E501
+    def forward(self,
+                input_ids: torch.Tensor,
+                positions: torch.Tensor,
+                kv_caches: List[torch.Tensor],
+                attn_metadata: AttentionMetadata,
+                image_input: Optional[torch.Tensor] = None) -> SamplerOutput:
         """Run forward pass for Llava 1.5.
 
         One key thing to understand is the `input_ids` already accounts for the
@@ -172,7 +174,7 @@ def forward(
                 image_features = image_input
             vision_embeddings = self.multi_modal_projector(image_features)
             inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-            _merge_vision_embeddings(
+            inputs_embeds = _merge_vision_embeddings(
                 input_ids, inputs_embeds, vision_embeddings,
                 self.vision_language_config.image_token_id)
             input_ids = None
diff --git a/vllm/model_executor/models/vlm_base.py b/vllm/model_executor/models/vlm_base.py
new file mode 100644
--- /dev/null
+++ b/vllm/model_executor/models/vlm_base.py
@@ -0,0 +1,12 @@
+from torch import nn
+
+from vllm.config import VisionLanguageConfig
+
+
+class VisionLanguageModelBase(nn.Module):
+    """Base class for all vision language models (VLMs)."""
+
+    def __init__(self, vision_language_config: VisionLanguageConfig) -> None:
+        super().__init__()
+
+        self.vision_language_config = vision_language_config


Test patch:
diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py
new file mode 100644
--- /dev/null
+++ b/tests/models/test_registry.py
@@ -0,0 +1,9 @@
+import pytest
+
+from vllm.model_executor.models import _MODELS, ModelRegistry
+
+
+@pytest.mark.parametrize("model_cls", _MODELS)
+def test_registry_imports(model_cls):
+    # Ensure all model classes can be imported successfully
+    ModelRegistry.load_model_cls(model_cls)