diff --git a/nemo_retriever/src/nemo_retriever/models/inference/vllm.py b/nemo_retriever/src/nemo_retriever/models/inference/vllm.py index fca2655e9e..6d84e53909 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/vllm.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/vllm.py @@ -16,12 +16,23 @@ from __future__ import annotations import logging +import os from typing import Any, List, Optional logger = logging.getLogger(__name__) VLLM_DTYPE = "bfloat16" VLLM_ATTENTION_BACKEND = "FLASH_ATTN" +VLLM_DEEP_GEMM_WARMUP_DEFAULT = "skip" + + +def apply_vllm_startup_defaults() -> None: + """Apply conservative vLLM startup defaults without overriding users.""" + + # DeepGEMM can still be used by vLLM at runtime. This only skips the + # ahead-of-time warmup path, which may fail before local inference starts + # when the optional DeepGEMM/CUDA-toolkit stack is not discoverable. + os.environ.setdefault("VLLM_DEEP_GEMM_WARMUP", VLLM_DEEP_GEMM_WARMUP_DEFAULT) def create_vllm_llm( @@ -44,6 +55,7 @@ def create_vllm_llm( Uses bfloat16 and FLASH_ATTN backend (fixed for this module). """ + apply_vllm_startup_defaults() try: from vllm import LLM except ImportError as e: @@ -185,4 +197,9 @@ def embed_multimodal_with_vllm_llm( return all_embeddings -__all__ = ["create_vllm_llm", "embed_with_vllm_llm", "embed_multimodal_with_vllm_llm"] +__all__ = [ + "apply_vllm_startup_defaults", + "create_vllm_llm", + "embed_with_vllm_llm", + "embed_multimodal_with_vllm_llm", +] diff --git a/nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py b/nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py index cfa833f4c9..c0b25450fc 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py +++ b/nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py @@ -81,6 +81,9 @@ def __init__( ) -> None: super().__init__() + from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults + + apply_vllm_startup_defaults() try: from vllm import LLM, SamplingParams # noqa: F401 except ImportError as e: diff --git a/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py b/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py index 14720fc95c..6eb7d0c1a9 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py +++ b/nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py @@ -74,6 +74,9 @@ def __init__( ) -> None: super().__init__() + from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults + + apply_vllm_startup_defaults() try: from vllm import LLM except ImportError as e: diff --git a/nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py b/nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py index 964546cf27..0f65932e51 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py +++ b/nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py @@ -84,6 +84,9 @@ def __init__( profile = get_caption_model_profile(model_path, target="local") model_path = profile.local_model_id + from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults + + apply_vllm_startup_defaults() try: from vllm import LLM, SamplingParams # noqa: F401 except ImportError as e: diff --git a/nemo_retriever/tests/test_vllm_embed.py b/nemo_retriever/tests/test_vllm_embed.py index 1ee99032e9..f71bb86874 100644 --- a/nemo_retriever/tests/test_vllm_embed.py +++ b/nemo_retriever/tests/test_vllm_embed.py @@ -6,6 +6,7 @@ import base64 import io +import os import sys from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock, patch @@ -14,7 +15,11 @@ torch = pytest.importorskip("torch") -from nemo_retriever.models.inference.vllm import embed_multimodal_with_vllm_llm, embed_with_vllm_llm +from nemo_retriever.models.inference.vllm import ( + apply_vllm_startup_defaults, + embed_multimodal_with_vllm_llm, + embed_with_vllm_llm, +) from nemo_retriever.models.local.llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2Embedder @@ -192,6 +197,32 @@ def test_limit_mm_per_prompt_forwarded_when_provided(self): _, kwargs = mock_llm_cls.call_args assert kwargs.get("limit_mm_per_prompt") == {"image": 1} + def test_applies_vllm_startup_defaults_before_constructing_llm(self, monkeypatch): + monkeypatch.delenv("VLLM_DEEP_GEMM_WARMUP", raising=False) + with patch("vllm.LLM") as mock_llm_cls: + mock_llm_cls.return_value = MagicMock() + from nemo_retriever.models.inference.vllm import create_vllm_llm + + create_vllm_llm("some-model") + + assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "skip" + + +class TestVllmStartupDefaults: + def test_deep_gemm_warmup_defaults_to_skip(self, monkeypatch): + monkeypatch.delenv("VLLM_DEEP_GEMM_WARMUP", raising=False) + + apply_vllm_startup_defaults() + + assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "skip" + + def test_deep_gemm_warmup_respects_user_override(self, monkeypatch): + monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "full") + + apply_vllm_startup_defaults() + + assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "full" + class TestVLLMEmbedderImages: def setup_method(self):