Preserve revision pins in secondary artifact loaders (#51446)

Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com>
Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com>

tests/config/test_model_arch_config.py

@@ -7,6 +7,7 @@ from copy import copy
 from pathlib import Path
 from types import SimpleNamespace
 from typing import cast
+from unittest.mock import Mock
 
 import pytest
 from transformers import PretrainedConfig
@@ -138,6 +139,32 @@ def test_head_size_falls_back_when_head_dim_is_zero():
     assert convertor.get_head_size() == 128
 
 
+def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.MonkeyPatch):
+    hf_config = PretrainedConfig(
+        model_type="deepseek_v2",
+        qk_rope_head_dim=128,
+        qk_nope_head_dim=128,
+    )
+    hf_config.name_or_path = "org/model"
+    get_hf_file_to_dict = Mock(return_value={"qk_rope_head_dim": 64})
+
+    monkeypatch.setattr(
+        "vllm.transformers_utils.repo_utils.get_hf_file_to_dict",
+        get_hf_file_to_dict,
+    )
+
+    convertor = ModelArchConfigConvertorBase(
+        hf_config,
+        hf_config,
+        revision="pinned-revision",
+    )
+
+    assert convertor._get_qk_rope_head_dim() == 64
+    get_hf_file_to_dict.assert_called_once_with(
+        "config.json", "org/model", "pinned-revision"
+    )
+
+
 @pytest.mark.parametrize(
     "attribute",
     [

tests/test_quantization_revision_pin.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm.config.vllm import VllmConfig
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+
+def test_quantization_config_forwards_model_revision(monkeypatch: pytest.MonkeyPatch):
+    quant_config = Mock()
+    quant_config.get_min_capability.return_value = 0
+    quant_config.get_supported_act_dtypes.return_value = [torch.float32]
+    model_config = SimpleNamespace(
+        quantization="auto_gptq",
+        model="repo/model",
+        hf_config=object(),
+        dtype=torch.float32,
+        revision="pinned-revision",
+    )
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.weight_utils.get_quant_config",
+        lambda model_config, load_config: quant_config,
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.current_platform.get_device_capability", lambda: None
+    )
+
+    VllmConfig._get_quantization_config(model_config, SimpleNamespace())
+
+    quant_config.maybe_update_config.assert_called_once_with(
+        "repo/model",
+        hf_config=model_config.hf_config,
+        revision="pinned-revision",
+    )

tests/test_tensorizer_revision_pin.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
+import io
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm.model_executor.model_loader import tensorizer
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+
+def test_serialize_vllm_model_forwards_model_revision(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    @contextlib.contextmanager
+    def fake_open_stream(*args, **kwargs):
+        yield io.BytesIO()
+
+    snapshot_download = Mock()
+    monkeypatch.setattr(tensorizer, "open_stream", fake_open_stream)
+    monkeypatch.setattr(tensorizer, "TensorSerializer", Mock())
+    monkeypatch.setattr(
+        tensorizer,
+        "hf_api",
+        lambda: SimpleNamespace(snapshot_download=snapshot_download),
+    )
+
+    model = torch.nn.Module()
+    tensorizer_config = SimpleNamespace(
+        _construct_tensorizer_args=lambda: SimpleNamespace(
+            tensorizer_uri="/tmp/model.tensors",
+            tensorizer_dir="/tmp/out",
+            stream_kwargs={},
+        ),
+        encryption_keyfile=None,
+        _is_sharded=False,
+        serialization_kwargs=None,
+    )
+    model_config = SimpleNamespace(
+        served_model_name="repo/model",
+        revision="pinned-revision",
+    )
+
+    tensorizer.serialize_vllm_model(model, tensorizer_config, model_config)
+
+    assert snapshot_download.call_args.args == ("repo/model",)
+    assert snapshot_download.call_args.kwargs["revision"] == "pinned-revision"

tests/transformers_utils/test_config.py

@@ -6,7 +6,15 @@ only get the `eos_token_id` from the tokenizer as defined by
 `BaseRenderer.get_eos_token_id`.
 """
 
+from types import SimpleNamespace
+from typing import cast
+from unittest.mock import patch
+
+from transformers import PretrainedConfig
+
+from vllm.config.model import ModelConfig
 from vllm.tokenizers import get_tokenizer
+from vllm.transformers_utils import config as config_module
 from vllm.transformers_utils.config import try_get_generation_config
 
 
@@ -30,3 +38,42 @@ def test_get_blip2_eos_token():
     generation_config = try_get_generation_config(model_name, trust_remote_code=False)
     assert generation_config is not None
     assert generation_config.eos_token_id == 50118
+
+
+def test_model_config_generation_fallback_forwards_code_revision():
+    model_config = cast(
+        ModelConfig,
+        SimpleNamespace(
+            generation_config="auto",
+            hf_config_path=None,
+            model="org/model",
+            trust_remote_code=True,
+            revision="model-pin",
+            code_revision="code-pin",
+            config_format="auto",
+            hf_token=None,
+        ),
+    )
+
+    with (
+        patch.object(
+            config_module.GenerationConfig,
+            "from_pretrained",
+            side_effect=OSError,
+        ),
+        patch.object(
+            config_module,
+            "get_config",
+            return_value=PretrainedConfig(),
+        ) as get_config,
+    ):
+        ModelConfig.try_get_generation_config(model_config)
+
+    get_config.assert_called_once_with(
+        "org/model",
+        trust_remote_code=True,
+        revision="model-pin",
+        code_revision="code-pin",
+        config_format="auto",
+        token=None,
+    )

vllm/config/model.py

@@ -890,7 +890,11 @@ class ModelConfig:
         convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(
             self.hf_config.model_type, ModelArchConfigConvertorBase
         )
-        convertor = convertor_cls(self.hf_config, self.hf_text_config)
+        convertor = convertor_cls(
+            self.hf_config,
+            self.hf_text_config,
+            revision=getattr(self, "revision", None),
+        )
         return convertor.convert(
             supports_multimodal=self._supports_multimodal_for_mm_prefix()
         )
@@ -1651,6 +1655,7 @@ class ModelConfig:
                 self.hf_config_path or self.model,
                 trust_remote_code=self.trust_remote_code,
                 revision=self.revision,
+                code_revision=self.code_revision,
                 config_format=self.config_format,
                 hf_token=self.hf_token,
             )
@@ -1658,6 +1663,7 @@ class ModelConfig:
             config = try_get_generation_config(
                 self.generation_config,
                 trust_remote_code=self.trust_remote_code,
+                code_revision=self.code_revision,
                 config_format=self.config_format,
                 hf_token=self.hf_token,
             )

vllm/config/vllm.py

@@ -779,6 +779,7 @@ class VllmConfig:
             quant_config.maybe_update_config(
                 model_config.model,
                 hf_config=model_config.hf_config,
+                revision=model_config.revision,
             )
             return quant_config
         return None

vllm/model_executor/model_loader/tensorizer.py

@@ -621,7 +621,9 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
 
 
 def serialize_extra_artifacts(
-    tensorizer_args: TensorizerArgs, served_model_name: str | list[str] | None
+    tensorizer_args: TensorizerArgs,
+    served_model_name: str | list[str] | None,
+    revision: str | None = None,
 ) -> None:
     if not isinstance(served_model_name, str):
         raise ValueError(
@@ -632,6 +634,7 @@ def serialize_extra_artifacts(
     with tempfile.TemporaryDirectory() as tmpdir:
         hf_api().snapshot_download(
             served_model_name,
+            revision=revision,
             local_dir=tmpdir,
             ignore_patterns=[
                 "*.pt",
@@ -693,7 +696,11 @@ def serialize_vllm_model(
         serializer.write_module(model)
         serializer.close()
 
-    serialize_extra_artifacts(tensorizer_args, model_config.served_model_name)
+    serialize_extra_artifacts(
+        tensorizer_args,
+        model_config.served_model_name,
+        revision=model_config.revision,
+    )
 
     logger.info("Successfully serialized model to %s", str(output_file))
     return model

vllm/model_executor/models/funaudiochat.py

@@ -519,12 +519,17 @@ class FunAudioChatProcessingInfo(BaseProcessingInfo):
 
     @cached_property
     def feature_extractor(self) -> WhisperFeatureExtractor:
-        return WhisperFeatureExtractor.from_pretrained(self.model_id)
+        return WhisperFeatureExtractor.from_pretrained(
+            self.model_id,
+            revision=self.ctx.model_config.revision,
+        )
 
     @cached_property
     def speech_tokenizer(self) -> TokenizersBackend:
         return TokenizersBackend.from_pretrained(
-            self.model_id, subfolder="speech_tokenizer"
+            self.model_id,
+            subfolder="speech_tokenizer",
+            revision=self.ctx.model_config.tokenizer_revision,
         )
 
     def get_feature_extractor(self) -> WhisperFeatureExtractor:

vllm/renderers/hf.py

@@ -210,7 +210,7 @@ def _build_mixed_prompt_embeds(
     return full_embeds, is_token_ids.tolist()
 
 
-_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
+_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, str | None, str | None, bool], str | None]()
 """
 Used in `_try_get_processor_chat_template` to avoid calling
 `cached_get_processor` again if the processor fails to be loaded.
@@ -222,9 +222,16 @@ This is needed because `lru_cache` does not cache when an exception happens.
 def _try_get_processor_chat_template(
     tokenizer: HfTokenizer,
     *,
+    revision: str | None,
+    code_revision: str | None,
     trust_remote_code: bool,
 ) -> str | None:
-    cache_key = (tokenizer.name_or_path, trust_remote_code)
+    cache_key = (
+        tokenizer.name_or_path,
+        revision,
+        code_revision,
+        trust_remote_code,
+    )
     if cache_key in _PROCESSOR_CHAT_TEMPLATES:
         return _PROCESSOR_CHAT_TEMPLATES[cache_key]
 
@@ -234,6 +241,8 @@ def _try_get_processor_chat_template(
         processor = cached_get_processor(
             tokenizer.name_or_path,
             processor_cls=(PythonBackend, TokenizersBackend, ProcessorMixin),
+            revision=revision,
+            code_revision=code_revision,
             trust_remote_code=trust_remote_code,
         )
         if (
@@ -271,6 +280,8 @@ def resolve_chat_template(
     if tools is None:
         chat_template = _try_get_processor_chat_template(
             tokenizer,
+            revision=model_config.revision,
+            code_revision=model_config.code_revision,
             trust_remote_code=model_config.trust_remote_code,
         )
         if chat_template is not None:

vllm/transformers_utils/config.py

@@ -1071,6 +1071,7 @@ def try_get_generation_config(
     model: str,
     trust_remote_code: bool,
     revision: str | None = None,
+    code_revision: str | None = None,
     config_format: str | ConfigFormat = "auto",
     hf_token: bool | str | None = None,
 ) -> GenerationConfig | None:
@@ -1086,6 +1087,7 @@ def try_get_generation_config(
                 model,
                 trust_remote_code=trust_remote_code,
                 revision=revision,
+                code_revision=code_revision,
                 config_format=config_format,
                 token=hf_token,
             )

vllm/transformers_utils/model_arch_config_convertor.py

@@ -24,9 +24,15 @@ logger = init_logger(__name__)
 
 
 class ModelArchConfigConvertorBase:
-    def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):
+    def __init__(
+        self,
+        hf_config: PretrainedConfig,
+        hf_text_config: PretrainedConfig,
+        revision: str | None = None,
+    ):
         self.hf_config = hf_config
         self.hf_text_config = hf_text_config
+        self.revision = revision
 
     def get_per_layer_hf_configs(
         self,
@@ -122,7 +128,7 @@ class ModelArchConfigConvertorBase:
         model_path = self.hf_config.name_or_path
         if not model_path:
             return qk_rope_head_dim
-        raw = get_hf_file_to_dict("config.json", model_path)
+        raw = get_hf_file_to_dict("config.json", model_path, self.revision)
         if raw and "qk_rope_head_dim" in raw:
             correct = raw["qk_rope_head_dim"]
             if correct != qk_rope_head_dim:
@@ -595,16 +601,26 @@ def _strip_mimo_v2_attention_chunk_size(
 
 
 class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase):
-    def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):
+    def __init__(
+        self,
+        hf_config: PretrainedConfig,
+        hf_text_config: PretrainedConfig,
+        revision: str | None = None,
+    ):
         if getattr(hf_config, "vision_config", None):
             hf_config.architectures = ["MiMoV2OmniForCausalLM"]
-        super().__init__(hf_config, hf_text_config)
+        super().__init__(hf_config, hf_text_config, revision)
         _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)
 
 
 class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
-    def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):
+    def __init__(
-        super().__init__(hf_config, hf_text_config)
+        self,
+        hf_config: PretrainedConfig,
+        hf_text_config: PretrainedConfig,
+        revision: str | None = None,
+    ):
+        super().__init__(hf_config, hf_text_config, revision)
         _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)
 
     def get_num_hidden_layers(self) -> int: