From 90fd4a333fd62b09310924dcb04d8e827bcbcd9b Mon Sep 17 00:00:00 2001 From: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:33:10 -0400 Subject: [PATCH] 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 | 27 ++++++++++ tests/test_quantization_revision_pin.py | 40 +++++++++++++++ tests/test_tensorizer_revision_pin.py | 51 +++++++++++++++++++ tests/transformers_utils/test_config.py | 47 +++++++++++++++++ vllm/config/model.py | 8 ++- vllm/config/vllm.py | 1 + .../model_executor/model_loader/tensorizer.py | 11 +++- vllm/model_executor/models/funaudiochat.py | 9 +++- vllm/renderers/hf.py | 15 +++++- vllm/transformers_utils/config.py | 2 + .../model_arch_config_convertor.py | 28 +++++++--- 11 files changed, 226 insertions(+), 13 deletions(-) create mode 100644 tests/test_quantization_revision_pin.py create mode 100644 tests/test_tensorizer_revision_pin.py diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index a4475de7cd..cf65c0e159 100644 --- a/tests/config/test_model_arch_config.py +++ b/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", [ diff --git a/tests/test_quantization_revision_pin.py b/tests/test_quantization_revision_pin.py new file mode 100644 index 0000000000..6c121ba031 --- /dev/null +++ b/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", + ) diff --git a/tests/test_tensorizer_revision_pin.py b/tests/test_tensorizer_revision_pin.py new file mode 100644 index 0000000000..5c2507afd1 --- /dev/null +++ b/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" diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 5a7421b6aa..42261feb53 100644 --- a/tests/transformers_utils/test_config.py +++ b/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, + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index e1d0e2568f..c09fdd7d76 100644 --- a/vllm/config/model.py +++ b/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, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 343feba36c..91372385e0 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index befe560929..3ef9b6bf2a 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/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 diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 2d3334f275..0a9e286783 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/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: diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 95210e4a47..6afd245cc1 100644 --- a/vllm/renderers/hf.py +++ b/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: diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 89da81ea91..ad8f0535b2 100644 --- a/vllm/transformers_utils/config.py +++ b/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, ) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 2e89c6aa32..2e6bd05cd9 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/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): - super().__init__(hf_config, hf_text_config) + def __init__( + 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: