[CI] Enable test_initialization to run on V1 (#16736)

Signed-off-by: mgoin <mgoin64@gmail.com>

tests/models/test_initialization.py

@@ -15,12 +15,12 @@ from .registry import HF_EXAMPLE_MODELS
 
 
 @pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())
-def test_can_initialize(model_arch):
+def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):
     model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)
     model_info.check_available_online(on_fail="skip")
     model_info.check_transformers_version(on_fail="skip")
 
-    # Avoid OOM
+    # Avoid OOM and reduce initialization time by only using 1 layer
     def hf_overrides(hf_config: PretrainedConfig) -> PretrainedConfig:
         hf_config.update(model_info.hf_overrides)
 
@@ -34,6 +34,12 @@ def test_can_initialize(model_arch):
             "num_local_experts": 2,
         })
 
+        if hasattr(hf_config, "vision_config"):
+            hf_config.vision_config.update({
+                "num_layers": 1,
+                "num_hidden_layers": 1,
+            })
+
         return hf_config
 
     # Avoid calling model.forward()
@@ -46,7 +52,7 @@ def test_can_initialize(model_arch):
         scheduler_kv_cache_config = get_kv_cache_config(
             vllm_config,
             kv_cache_specs[0],
-            20 * GiB_bytes,
+            10 * GiB_bytes,
         )
 
         # gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
@@ -55,7 +61,9 @@ def test_can_initialize(model_arch):
     with (patch.object(V0LLMEngine, "_initialize_kv_caches",
                        _initialize_kv_caches_v0),
           patch.object(V1EngineCore, "_initialize_kv_caches",
-                       _initialize_kv_caches_v1)):
+                       _initialize_kv_caches_v1), monkeypatch.context() as m):
+        if model_info.v0_only:
+            m.setenv("VLLM_USE_V1", "0")
         LLM(
             model_info.default,
             tokenizer=model_info.tokenizer,
@@ -65,6 +73,7 @@ def test_can_initialize(model_arch):
                 "num_speculative_tokens": 1,
             } if model_info.speculative_model else None,
             trust_remote_code=model_info.trust_remote_code,
+            max_model_len=model_info.max_model_len,
             load_format="dummy",
             hf_overrides=hf_overrides,
         )