90fd4a333f · tduncan/Nvidia-Optimization-Stack-VLLM

tests/config/test_model_arch_config.py

+27

@@ -7,6 +7,7 @@ from copy import copy

7 from pathlib import Path

8 from types import SimpleNamespace

9 from typing import cast

10 +from unittest.mock import Mock

11

12 import pytest

13 from transformers import PretrainedConfig

@@ -138,6 +139,32 @@ def test_head_size_falls_back_when_head_dim_is_zero():

139 assert convertor.get_head_size() == 128

140

141

142 +def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.MonkeyPatch):

143 + hf_config = PretrainedConfig(

144 + model_type="deepseek_v2",

145 + qk_rope_head_dim=128,

146 + qk_nope_head_dim=128,

147 + )

148 + hf_config.name_or_path = "org/model"

149 + get_hf_file_to_dict = Mock(return_value={"qk_rope_head_dim": 64})

150 +

151 + monkeypatch.setattr(

152 + "vllm.transformers_utils.repo_utils.get_hf_file_to_dict",

153 + get_hf_file_to_dict,

154 + )

155 +

156 + convertor = ModelArchConfigConvertorBase(

157 + hf_config,

158 + hf_config,

159 + revision="pinned-revision",

160 + )

161 +

162 + assert convertor._get_qk_rope_head_dim() == 64

163 + get_hf_file_to_dict.assert_called_once_with(

164 + "config.json", "org/model", "pinned-revision"

165 + )

166 +

167 +

168 @pytest.mark.parametrize(

169 "attribute",

170 [

tests/test_quantization_revision_pin.py new

+40

@@ -0,0 +1,40 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from types import SimpleNamespace

4 +from unittest.mock import Mock

5 +

6 +import pytest

7 +import torch

8 +

9 +from vllm.config.vllm import VllmConfig

10 +

11 +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]

12 +

13 +

14 +def test_quantization_config_forwards_model_revision(monkeypatch: pytest.MonkeyPatch):

15 + quant_config = Mock()

16 + quant_config.get_min_capability.return_value = 0

17 + quant_config.get_supported_act_dtypes.return_value = [torch.float32]

18 + model_config = SimpleNamespace(

19 + quantization="auto_gptq",

20 + model="repo/model",

21 + hf_config=object(),

22 + dtype=torch.float32,

23 + revision="pinned-revision",

24 + )

25 +

26 + monkeypatch.setattr(

27 + "vllm.model_executor.model_loader.weight_utils.get_quant_config",

28 + lambda model_config, load_config: quant_config,

29 + )

30 + monkeypatch.setattr(

31 + "vllm.platforms.current_platform.get_device_capability", lambda: None

32 + )

33 +

34 + VllmConfig._get_quantization_config(model_config, SimpleNamespace())

35 +

36 + quant_config.maybe_update_config.assert_called_once_with(

37 + "repo/model",

38 + hf_config=model_config.hf_config,

39 + revision="pinned-revision",

40 + )

tests/test_tensorizer_revision_pin.py new

+51

@@ -0,0 +1,51 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import contextlib

4 +import io

5 +from types import SimpleNamespace

6 +from unittest.mock import Mock

7 +

8 +import pytest

9 +import torch

10 +

11 +from vllm.model_executor.model_loader import tensorizer

12 +

13 +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]

14 +

15 +

16 +def test_serialize_vllm_model_forwards_model_revision(

17 + monkeypatch: pytest.MonkeyPatch,

18 +):

19 + @contextlib.contextmanager

20 + def fake_open_stream(*args, **kwargs):

21 + yield io.BytesIO()

22 +

23 + snapshot_download = Mock()

24 + monkeypatch.setattr(tensorizer, "open_stream", fake_open_stream)

25 + monkeypatch.setattr(tensorizer, "TensorSerializer", Mock())

26 + monkeypatch.setattr(

27 + tensorizer,

28 + "hf_api",

29 + lambda: SimpleNamespace(snapshot_download=snapshot_download),

30 + )

31 +

32 + model = torch.nn.Module()

33 + tensorizer_config = SimpleNamespace(

34 + _construct_tensorizer_args=lambda: SimpleNamespace(

35 + tensorizer_uri="/tmp/model.tensors",

36 + tensorizer_dir="/tmp/out",

37 + stream_kwargs={},

38 + ),

39 + encryption_keyfile=None,

40 + _is_sharded=False,

41 + serialization_kwargs=None,

42 + )

43 + model_config = SimpleNamespace(

44 + served_model_name="repo/model",

45 + revision="pinned-revision",

46 + )

47 +

48 + tensorizer.serialize_vllm_model(model, tensorizer_config, model_config)

49 +

50 + assert snapshot_download.call_args.args == ("repo/model",)

51 + assert snapshot_download.call_args.kwargs["revision"] == "pinned-revision"

tests/transformers_utils/test_config.py

+47

@@ -6,7 +6,15 @@ only get the `eos_token_id` from the tokenizer as defined by

6 `BaseRenderer.get_eos_token_id`.

7 """

8

9 +from types import SimpleNamespace

10 +from typing import cast

11 +from unittest.mock import patch

12 +

13 +from transformers import PretrainedConfig

14 +

15 +from vllm.config.model import ModelConfig

16 from vllm.tokenizers import get_tokenizer

17 +from vllm.transformers_utils import config as config_module

18 from vllm.transformers_utils.config import try_get_generation_config

19

20

@@ -30,3 +38,42 @@ def test_get_blip2_eos_token():

38 generation_config = try_get_generation_config(model_name, trust_remote_code=False)

39 assert generation_config is not None

40 assert generation_config.eos_token_id == 50118

41 +

42 +

43 +def test_model_config_generation_fallback_forwards_code_revision():

44 + model_config = cast(

45 + ModelConfig,

46 + SimpleNamespace(

47 + generation_config="auto",

48 + hf_config_path=None,

49 + model="org/model",

50 + trust_remote_code=True,

51 + revision="model-pin",

52 + code_revision="code-pin",

53 + config_format="auto",

54 + hf_token=None,

55 + ),

56 + )

57 +

58 + with (

59 + patch.object(

60 + config_module.GenerationConfig,

61 + "from_pretrained",

62 + side_effect=OSError,

63 + ),

64 + patch.object(

65 + config_module,

66 + "get_config",

67 + return_value=PretrainedConfig(),

68 + ) as get_config,

69 + ):

70 + ModelConfig.try_get_generation_config(model_config)

71 +

72 + get_config.assert_called_once_with(

73 + "org/model",

74 + trust_remote_code=True,

75 + revision="model-pin",

76 + code_revision="code-pin",

77 + config_format="auto",

78 + token=None,

79 + )

vllm/config/model.py

+7 -1

@@ -890,7 +890,11 @@ class ModelConfig:

890 convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(

891 self.hf_config.model_type, ModelArchConfigConvertorBase

892 )

893 - convertor = convertor_cls(self.hf_config, self.hf_text_config)

893 + convertor = convertor_cls(

894 + self.hf_config,

895 + self.hf_text_config,

896 + revision=getattr(self, "revision", None),

897 + )

898 return convertor.convert(

899 supports_multimodal=self._supports_multimodal_for_mm_prefix()

900 )

@@ -1651,6 +1655,7 @@ class ModelConfig:

1655 self.hf_config_path or self.model,

1656 trust_remote_code=self.trust_remote_code,

1657 revision=self.revision,

1658 + code_revision=self.code_revision,

1659 config_format=self.config_format,

1660 hf_token=self.hf_token,

1661 )

@@ -1658,6 +1663,7 @@ class ModelConfig:

1663 config = try_get_generation_config(

1664 self.generation_config,

1665 trust_remote_code=self.trust_remote_code,

1666 + code_revision=self.code_revision,

1667 config_format=self.config_format,

1668 hf_token=self.hf_token,

1669 )

vllm/config/vllm.py

+1

@@ -779,6 +779,7 @@ class VllmConfig:

779 quant_config.maybe_update_config(

780 model_config.model,

781 hf_config=model_config.hf_config,

782 + revision=model_config.revision,

783 )

784 return quant_config

785 return None

vllm/model_executor/model_loader/tensorizer.py

+9 -2

@@ -621,7 +621,9 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:

621

622

623 def serialize_extra_artifacts(

624 - tensorizer_args: TensorizerArgs, served_model_name: str | list[str] | None

624 + tensorizer_args: TensorizerArgs,

625 + served_model_name: str | list[str] | None,

626 + revision: str | None = None,

627 ) -> None:

628 if not isinstance(served_model_name, str):

629 raise ValueError(

@@ -632,6 +634,7 @@ def serialize_extra_artifacts(

634 with tempfile.TemporaryDirectory() as tmpdir:

635 hf_api().snapshot_download(

636 served_model_name,

637 + revision=revision,

638 local_dir=tmpdir,

639 ignore_patterns=[

640 "*.pt",

@@ -693,7 +696,11 @@ def serialize_vllm_model(

696 serializer.write_module(model)

697 serializer.close()

698

696 - serialize_extra_artifacts(tensorizer_args, model_config.served_model_name)

699 + serialize_extra_artifacts(

700 + tensorizer_args,

701 + model_config.served_model_name,

702 + revision=model_config.revision,

703 + )

704

705 logger.info("Successfully serialized model to %s", str(output_file))

706 return model

vllm/model_executor/models/funaudiochat.py

+7 -2

@@ -519,12 +519,17 @@ class FunAudioChatProcessingInfo(BaseProcessingInfo):

519

520 @cached_property

521 def feature_extractor(self) -> WhisperFeatureExtractor:

522 - return WhisperFeatureExtractor.from_pretrained(self.model_id)

522 + return WhisperFeatureExtractor.from_pretrained(

523 + self.model_id,

524 + revision=self.ctx.model_config.revision,

525 + )

526

527 @cached_property

528 def speech_tokenizer(self) -> TokenizersBackend:

529 return TokenizersBackend.from_pretrained(

527 - self.model_id, subfolder="speech_tokenizer"

530 + self.model_id,

531 + subfolder="speech_tokenizer",

532 + revision=self.ctx.model_config.tokenizer_revision,

533 )

534

535 def get_feature_extractor(self) -> WhisperFeatureExtractor:

vllm/renderers/hf.py

+13 -2

@@ -210,7 +210,7 @@ def _build_mixed_prompt_embeds(

210 return full_embeds, is_token_ids.tolist()

211

212

213 -_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()

213 +_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, str | None, str | None, bool], str | None]()

214 """

215 Used in `_try_get_processor_chat_template` to avoid calling

216 `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.

222 def _try_get_processor_chat_template(

223 tokenizer: HfTokenizer,

224 *,

225 + revision: str | None,

226 + code_revision: str | None,

227 trust_remote_code: bool,

228 ) -> str | None:

227 - cache_key = (tokenizer.name_or_path, trust_remote_code)

229 + cache_key = (

230 + tokenizer.name_or_path,

231 + revision,

232 + code_revision,

233 + trust_remote_code,

234 + )

235 if cache_key in _PROCESSOR_CHAT_TEMPLATES:

236 return _PROCESSOR_CHAT_TEMPLATES[cache_key]

237

@@ -234,6 +241,8 @@ def _try_get_processor_chat_template(

241 processor = cached_get_processor(

242 tokenizer.name_or_path,

243 processor_cls=(PythonBackend, TokenizersBackend, ProcessorMixin),

244 + revision=revision,

245 + code_revision=code_revision,

246 trust_remote_code=trust_remote_code,

247 )

248 if (

@@ -271,6 +280,8 @@ def resolve_chat_template(

280 if tools is None:

281 chat_template = _try_get_processor_chat_template(

282 tokenizer,

283 + revision=model_config.revision,

284 + code_revision=model_config.code_revision,

285 trust_remote_code=model_config.trust_remote_code,

286 )

287 if chat_template is not None:

vllm/transformers_utils/config.py

+2

@@ -1071,6 +1071,7 @@ def try_get_generation_config(

1071 model: str,

1072 trust_remote_code: bool,

1073 revision: str | None = None,

1074 + code_revision: str | None = None,

1075 config_format: str | ConfigFormat = "auto",

1076 hf_token: bool | str | None = None,

1077 ) -> GenerationConfig | None:

@@ -1086,6 +1087,7 @@ def try_get_generation_config(

1087 model,

1088 trust_remote_code=trust_remote_code,

1089 revision=revision,

1090 + code_revision=code_revision,

1091 config_format=config_format,

1092 token=hf_token,

1093 )

vllm/transformers_utils/model_arch_config_convertor.py

+22 -6

@@ -24,9 +24,15 @@ logger = init_logger(__name__)

24

25

26 class ModelArchConfigConvertorBase:

27 - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):

27 + def __init__(

28 + self,

29 + hf_config: PretrainedConfig,

30 + hf_text_config: PretrainedConfig,

31 + revision: str | None = None,

32 + ):

33 self.hf_config = hf_config

34 self.hf_text_config = hf_text_config

35 + self.revision = revision

36

37 def get_per_layer_hf_configs(

38 self,

@@ -122,7 +128,7 @@ class ModelArchConfigConvertorBase:

128 model_path = self.hf_config.name_or_path

129 if not model_path:

130 return qk_rope_head_dim

125 - raw = get_hf_file_to_dict("config.json", model_path)

131 + raw = get_hf_file_to_dict("config.json", model_path, self.revision)

132 if raw and "qk_rope_head_dim" in raw:

133 correct = raw["qk_rope_head_dim"]

134 if correct != qk_rope_head_dim:

@@ -595,16 +601,26 @@ def _strip_mimo_v2_attention_chunk_size(

601

602

603 class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase):

598 - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):

604 + def __init__(

605 + self,

606 + hf_config: PretrainedConfig,

607 + hf_text_config: PretrainedConfig,

608 + revision: str | None = None,

609 + ):

610 if getattr(hf_config, "vision_config", None):

611 hf_config.architectures = ["MiMoV2OmniForCausalLM"]

601 - super().__init__(hf_config, hf_text_config)

612 + super().__init__(hf_config, hf_text_config, revision)

613 _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)

614

615

616 class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

606 - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):

607 - super().__init__(hf_config, hf_text_config)

617 + def __init__(

618 + self,

619 + hf_config: PretrainedConfig,

620 + hf_text_config: PretrainedConfig,

621 + revision: str | None = None,

622 + ):

623 + super().__init__(hf_config, hf_text_config, revision)

624 _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)

625

626 def get_num_hidden_layers(self) -> int:

Preserve revision pins in secondary artifact loaders (#51446)