@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6aa8f9a4e7

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

[Core] Rework dtype resolution (#18751)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jun 1, 2025 at 11:04 UTC 6aa8f9a4e7ed7af459476afe4f293f383e35f3e8

13 files changed +314 -119

tests/basic_correctness/test_basic_correctness.py

+1 -4

@@ -60,7 +60,6 @@ def _fix_prompt_embed_outputs(

61 @pytest.mark.parametrize("model", MODELS)

62 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])

63 -@pytest.mark.parametrize("dtype", ["half"])

63 @pytest.mark.parametrize("max_tokens", [5])

64 @pytest.mark.parametrize("enforce_eager", [False])

65 @pytest.mark.parametrize("enable_prompt_embeds", [True, False])

@@ -69,7 +68,6 @@ def test_models(

68 hf_runner,

69 model: str,

70 backend: str,

72 - dtype: str,

71 max_tokens: int,

72 enforce_eager: bool,

73 enable_prompt_embeds: bool,

@@ -97,7 +95,7 @@ def test_models(

95 str(i) for i in range(1024)) + " are:"

96 example_prompts = [prompt]

100 - with hf_runner(model, dtype=dtype) as hf_model:

98 + with hf_runner(model) as hf_model:

99 hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)

100 if enable_prompt_embeds:

101 with torch.no_grad():

@@ -106,7 +104,6 @@ def test_models(

104

105 with VllmRunner(model,

106 max_model_len=8192,

109 - dtype=dtype,

107 enforce_eager=enforce_eager,

108 enable_prompt_embeds=enable_prompt_embeds,

109 gpu_memory_utilization=0.7) as vllm_model:

tests/conftest.py

+6 -1

@@ -324,7 +324,12 @@ class HfRunner:

324 trust_remote_code=trust_remote_code,

325 )

326 self.device = self.get_default_device()

327 - self.dtype = torch_dtype = _get_and_verify_dtype(self.config, dtype)

327 + self.dtype = torch_dtype = _get_and_verify_dtype(

328 + self.model_name,

329 + self.config,

330 + dtype=dtype,

331 + is_pooling_model=is_sentence_transformer or is_cross_encoder,

332 + )

333

334 model_kwargs = model_kwargs if model_kwargs is not None else {}

335 model_kwargs.setdefault("torch_dtype", torch_dtype)

tests/models/language/pooling/mteb_utils.py

+4 -7

@@ -102,21 +102,18 @@ def mteb_test_embed_models(hf_runner,

102 vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),

103 MTEB_EMBED_TASKS)

104 vllm_dtype = vllm_model.model.llm_engine.model_config.dtype

105 - model_dtype = getattr(

106 - vllm_model.model.llm_engine.model_config.hf_config, "torch_dtype",

107 - vllm_dtype)

105

109 - with set_default_torch_dtype(model_dtype) and hf_runner(

106 + with set_default_torch_dtype(vllm_dtype) and hf_runner(

107 model_info.name, is_sentence_transformer=True,

111 - dtype=model_dtype) as hf_model:

108 + dtype=vllm_dtype) as hf_model:

109

110 if hf_model_callback is not None:

111 hf_model_callback(hf_model)

112

113 st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)

114

118 - print("VLLM:", vllm_dtype, vllm_main_score)

119 - print("SentenceTransformer:", model_dtype, st_main_score)

115 + print("VLLM:", vllm_main_score)

116 + print("SentenceTransformers:", st_main_score)

117 print("Difference:", st_main_score - vllm_main_score)

118

119 assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_EMBED_TOL)

tests/models/language/pooling/test_classification.py

+1 -1

@@ -43,6 +43,6 @@ def test_models(

44 # the tolerance value of 1e-2 is selected based on the

45 # half datatype tests in

46 - # tests/models/embedding/language/test_embedding.py

46 + # tests/models/language/pooling/test_embedding.py

47 assert torch.allclose(hf_output, vllm_output,

48 1e-3 if dtype == "float" else 1e-2)

tests/models/language/pooling/test_embedding.py

+1 -5

@@ -30,13 +30,11 @@ from ...utils import check_embeddings_close

30 pytest.param("sentence-transformers/stsb-roberta-base-v2"),

31 ],

32 )

33 -@pytest.mark.parametrize("dtype", ["half"])

33 def test_models(

34 hf_runner,

35 vllm_runner,

36 example_prompts,

37 model,

39 - dtype: str,

38 monkeypatch,

39 ) -> None:

@@ -58,13 +56,11 @@ def test_models(

56 # So we need to strip the input texts to avoid test failing.

57 example_prompts = [str(s).strip() for s in example_prompts]

61 - with hf_runner(model, dtype=dtype,

62 - is_sentence_transformer=True) as hf_model:

59 + with hf_runner(model, is_sentence_transformer=True) as hf_model:

60 hf_outputs = hf_model.encode(example_prompts)

62 with vllm_runner(model,

63 task="embed",

67 - dtype=dtype,

64 max_model_len=None,

65 **vllm_extra_kwargs) as vllm_model:

66 vllm_outputs = vllm_model.encode(example_prompts)

tests/models/multimodal/generation/test_whisper.py

@@ -100,6 +100,7 @@ def run_test(

100

101 with vllm_runner(

102 model,

103 + dtype="half",

104 max_model_len=448,

105 tensor_parallel_size=tensor_parallel_size,

106 distributed_executor_backend=distributed_executor_backend,

tests/models/multimodal/processing/test_common.py

+1 -1

@@ -40,7 +40,7 @@ def _test_processing_correctness(

40 tokenizer_mode=model_info.tokenizer_mode,

41 trust_remote_code=model_info.trust_remote_code,

42 seed=0,

43 - dtype="float16",

43 + dtype="auto",

44 revision=None,

45 hf_overrides=model_info.hf_overrides,

46 )

tests/samplers/test_no_bad_words.py

+1 -1

@@ -103,7 +103,7 @@ class TestTwoTokenBadWord:

103 add_special_tokens=False)[0]

104

105 def test_two_token_bad_word(self, vllm_runner):

106 - with vllm_runner(self.MODEL) as llm:

106 + with vllm_runner(self.MODEL, dtype="half") as llm:

107 output_token_ids = self._generate(llm)

108 assert output_token_ids[:2] == [

109 self.target_token_id1, self.target_token_id2

tests/test_utils.py

+80 -22

@@ -17,7 +17,8 @@ from vllm_test_utils.monitor import monitor

17 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

18 from vllm.utils import (CacheInfo, FlexibleArgumentParser, LRUCache,

19 MemorySnapshot, PlaceholderModule, StoreBoolean,

20 - bind_kv_cache, deprecate_kwargs, get_open_port,

20 + bind_kv_cache, common_broadcastable_dtype,

21 + deprecate_kwargs, get_open_port, is_lossless_cast,

22 make_zmq_path, make_zmq_socket, memory_profiling,

23 merge_async_iterators, sha256, split_zmq_path,

24 supports_kw, swap_dict_values)

@@ -567,12 +568,65 @@ def test_lru_cache():

568 assert 6 in cache

569

570

571 +# yapf: disable

572 +@pytest.mark.parametrize(

573 + ("src_dtype", "tgt_dtype", "expected_result"),

574 + [

575 + # Different precision_levels

576 + (torch.bool, torch.int8, True),

577 + (torch.bool, torch.float16, True),

578 + (torch.bool, torch.complex32, True),

579 + (torch.int64, torch.bool, False),

580 + (torch.int64, torch.float16, True),

581 + (torch.int64, torch.complex32, True),

582 + (torch.float64, torch.bool, False),

583 + (torch.float64, torch.int8, False),

584 + (torch.float64, torch.complex32, True),

585 + (torch.complex128, torch.bool, False),

586 + (torch.complex128, torch.int8, False),

587 + (torch.complex128, torch.float16, False),

588 + # precision_level=0

589 + (torch.bool, torch.bool, True),

590 + # precision_level=1

591 + (torch.int8, torch.int16, True),

592 + (torch.int16, torch.int8, False),

593 + (torch.uint8, torch.int8, False),

594 + (torch.int8, torch.uint8, False),

595 + # precision_level=2

596 + (torch.float16, torch.float32, True),

597 + (torch.float32, torch.float16, False),

598 + (torch.bfloat16, torch.float32, True),

599 + (torch.float32, torch.bfloat16, False),

600 + # precision_level=3

601 + (torch.complex32, torch.complex64, True),

602 + (torch.complex64, torch.complex32, False),

603 + ],

604 +)

605 +# yapf: enable

606 +def test_is_lossless_cast(src_dtype, tgt_dtype, expected_result):

607 + assert is_lossless_cast(src_dtype, tgt_dtype) == expected_result

608 +

609 +

610 +# yapf: disable

611 +@pytest.mark.parametrize(

612 + ("dtypes", "expected_result"),

613 + [

614 + ([torch.bool], torch.bool),

615 + ([torch.bool, torch.int8], torch.int8),

616 + ([torch.bool, torch.int8, torch.float16], torch.float16),

617 + ([torch.bool, torch.int8, torch.float16, torch.complex32], torch.complex32), # noqa: E501

618 + ],

619 +)

620 +# yapf: enable

621 +def test_common_broadcastable_dtype(dtypes, expected_result):

622 + assert common_broadcastable_dtype(dtypes) == expected_result

623 +

624 +

625 def test_placeholder_module_error_handling():

626 placeholder = PlaceholderModule("placeholder_1234")

627

628 def build_ctx():

574 - return pytest.raises(ModuleNotFoundError,

575 - match="No module named")

629 + return pytest.raises(ModuleNotFoundError, match="No module named")

630

631 with build_ctx():

632 int(placeholder)

@@ -608,6 +662,7 @@ def test_placeholder_module_error_handling():

662 _ = placeholder_attr.module

663

664

665 +# yapf: disable

666 @pytest.mark.parametrize(

667 "obj,key1,key2",

668 [

@@ -618,6 +673,7 @@ def test_placeholder_module_error_handling():

673 # Tests for both keys do not exist

674 ({1: "a", 2: "b"}, 3, 4),

675 ])

676 +# yapf: enable

677 def test_swap_dict_values(obj, key1, key2):

678 original_obj = obj.copy()

679 swap_dict_values(obj, key1, key2)

@@ -631,19 +687,19 @@ def test_swap_dict_values(obj, key1, key2):

687 assert key1 not in obj

688

689

634 -def test_model_specification(parser_with_config,

635 - cli_config_file,

690 +def test_model_specification(parser_with_config, cli_config_file,

691 cli_config_file_with_model):

692 # Test model in CLI takes precedence over config

638 - args = parser_with_config.parse_args([

639 - 'serve', 'cli-model', '--config', cli_config_file_with_model

640 - ])

693 + args = parser_with_config.parse_args(

694 + ['serve', 'cli-model', '--config', cli_config_file_with_model])

695 assert args.model_tag == 'cli-model'

696 assert args.served_model_name == 'mymodel'

697

698 # Test model from config file works

699 args = parser_with_config.parse_args([

646 - 'serve', '--config', cli_config_file_with_model,

700 + 'serve',

701 + '--config',

702 + cli_config_file_with_model,

703 ])

704 assert args.model == 'config-model'

705 assert args.served_model_name == 'mymodel'

@@ -654,17 +710,19 @@ def test_model_specification(parser_with_config,

710

711 # Test using --model option raises error

712 with pytest.raises(

657 - ValueError,

658 - match=(

659 - "With `vllm serve`, you should provide the model as a positional "

660 - "argument or in a config file instead of via the `--model` option."

661 - ),

713 + ValueError,

714 + match=

715 + ("With `vllm serve`, you should provide the model as a positional "

716 + "argument or in a config file instead of via the `--model` option."),

717 ):

718 parser_with_config.parse_args(['serve', '--model', 'my-model'])

719

720 # Test other config values are preserved

721 args = parser_with_config.parse_args([

667 - 'serve', 'cli-model', '--config', cli_config_file_with_model,

722 + 'serve',

723 + 'cli-model',

724 + '--config',

725 + cli_config_file_with_model,

726 ])

727 assert args.tensor_parallel_size == 2

728 assert args.trust_remote_code is True

@@ -673,7 +731,7 @@ def test_model_specification(parser_with_config,

731

732

733 @pytest.mark.parametrize("input", [(), ("abc", ), (None, ),

676 - (None, bool, [1, 2, 3])])

734 + (None, bool, [1, 2, 3])])

735 @pytest.mark.parametrize("output", [0, 1, 2])

736 def test_sha256(input: tuple, output: int):

737 hash = sha256(input)

@@ -682,7 +740,8 @@ def test_sha256(input: tuple, output: int):

740 assert hash != 0

741

742 bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)

685 - assert hash == int.from_bytes(hashlib.sha256(bytes).digest(), byteorder="big")

743 + assert hash == int.from_bytes(hashlib.sha256(bytes).digest(),

744 + byteorder="big")

745

746 # hashing again, returns the same value

747 assert hash == sha256(input)

@@ -698,8 +757,7 @@ def test_sha256(input: tuple, output: int):

757 ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),

758 ("tcp://[::1]:5555", ("tcp", "::1", "5555")), # IPv6 address

759 ("inproc://some_identifier", ("inproc", "some_identifier", "")),

701 - ]

702 -)

760 + ])

761 def test_split_zmq_path(path, expected):

762 assert split_zmq_path(path) == expected

763

@@ -711,8 +769,7 @@ def test_split_zmq_path(path, expected):

769 "tcp://127.0.0.1", # Missing port

770 "tcp://[::1]", # Missing port for IPv6

771 "tcp://:5555", # Missing host

714 - ]

715 -)

772 + ])

773 def test_split_zmq_path_invalid(invalid_path):

774 with pytest.raises(ValueError):

775 split_zmq_path(invalid_path)

@@ -734,7 +791,8 @@ def test_make_zmq_socket_ipv6():

791 zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)

792

793 # Verify that the IPV6 option is set

737 - assert zsock.getsockopt(zmq.IPV6) == 1, "IPV6 option should be enabled for IPv6 addresses"

794 + assert zsock.getsockopt(

795 + zmq.IPV6) == 1, "IPV6 option should be enabled for IPv6 addresses"

796

797 # Clean up

798 zsock.close()

vllm/config.py

+135 -67

@@ -24,6 +24,7 @@ import torch

24 from pydantic import (ConfigDict, SkipValidation, TypeAdapter, field_validator,

25 model_validator)

26 from pydantic.dataclasses import dataclass

27 +from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

28 from torch.distributed import ProcessGroup, ReduceOp

29 from transformers import PretrainedConfig

30 from typing_extensions import deprecated, runtime_checkable

@@ -42,15 +43,16 @@ from vllm.transformers_utils.config import (

43 ConfigFormat, get_config, get_hf_image_processor_config,

44 get_hf_text_config, get_pooling_config,

45 get_sentence_transformer_tokenizer_config, is_encoder_decoder,

45 - try_get_generation_config, uses_mrope)

46 + try_get_generation_config, try_get_safetensors_metadata, uses_mrope)

47 from vllm.transformers_utils.s3_utils import S3Model

48 from vllm.transformers_utils.utils import is_s3, maybe_model_redirect

49 from vllm.utils import (DEFAULT_MAX_NUM_BATCHED_TOKENS,

50 MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,

51 POOLING_MODEL_MAX_NUM_BATCHED_TOKENS, GiB_bytes,

51 - LayerBlockType, cuda_device_count_stateless,

52 - get_cpu_memory, get_open_port, is_torch_equal_or_newer,

53 - random_uuid, resolve_obj_by_qualname)

52 + LayerBlockType, common_broadcastable_dtype,

53 + cuda_device_count_stateless, get_cpu_memory,

54 + get_open_port, is_torch_equal_or_newer, random_uuid,

55 + resolve_obj_by_qualname)

57 if TYPE_CHECKING:

58 from _typeshed import DataclassInstance

@@ -540,7 +542,24 @@ class ModelConfig:

542 self.encoder_config = self._get_encoder_config()

543 self.hf_image_processor_config = get_hf_image_processor_config(

544 self.model, hf_token=self.hf_token, revision=self.revision)

543 - self.dtype = _get_and_verify_dtype(self.hf_config, self.dtype)

545 +

546 + supported_tasks, task = self._resolve_task(self.task)

547 + self.supported_tasks = supported_tasks

548 + self.task = task

549 + if self.task in ("draft", "generate"):

550 + self.truncation_side = "left"

551 + else:

552 + self.truncation_side = "right"

553 +

554 + self.pooler_config = self._init_pooler_config()

555 +

556 + self.dtype = _get_and_verify_dtype(

557 + self.model,

558 + self.hf_config,

559 + self.dtype,

560 + is_pooling_model=self.runner_type == "pooling",

561 + revision=self.revision,

562 + )

563

564 # Workaround for Gemma 2 which uses interleaved sliding window

565 # attention, but it's not specified in its config. TODO: remove this

@@ -597,16 +616,6 @@ class ModelConfig:

616 raise ValueError(

617 "`override_neuron_config` is only supported on Neuron.")

618

600 - supported_tasks, task = self._resolve_task(self.task)

601 - self.supported_tasks = supported_tasks

602 - self.task = task

603 - if self.task in ("draft", "generate"):

604 - self.truncation_side = "left"

605 - else:

606 - self.truncation_side = "right"

607 -

608 - self.pooler_config = self._init_pooler_config()

609 -

619 self._verify_quantization()

620 self._verify_cuda_graph()

621 self._verify_bnb_config()

@@ -692,7 +701,6 @@ class ModelConfig:

701 self.model, self.revision)

702

703 def _init_pooler_config(self) -> Optional["PoolerConfig"]:

695 -

704 if self.runner_type == "pooling":

705 if isinstance(self.override_pooler_config, dict):

706 self.override_pooler_config = PoolerConfig(

@@ -3074,13 +3082,37 @@ _STR_DTYPE_TO_TORCH_DTYPE = {

3082 "bfloat16": torch.bfloat16,

3083 }

3084

3077 -_ROCM_NOT_SUPPORTED_DTYPE: list[str] = [] #

3085 +# model_type -> reason

3086 +_FLOAT16_NOT_SUPPORTED_MODELS = {

3087 + "gemma2": "Numerical instability. Please use bfloat16 or float32 instead.",

3088 + "gemma3": "Numerical instability. Please use bfloat16 or float32 instead.",

3089 + "plamo2": "Numerical instability. Please use bfloat16 or float32 instead.",

3090 + "glm4": "Numerical instability. Please use bfloat16 or float32 instead.",

3091 +}

3092

3093

3080 -def _get_and_verify_dtype(

3094 +def _is_valid_dtype(model_type: str, dtype: torch.dtype):

3095 + if model_type in _FLOAT16_NOT_SUPPORTED_MODELS and dtype == torch.float16: # noqa: E501, SIM103

3096 + return False

3097 +

3098 + return True

3099 +

3100 +

3101 +def _check_valid_dtype(model_type: str, dtype: torch.dtype):

3102 + if model_type in _FLOAT16_NOT_SUPPORTED_MODELS and dtype == torch.float16:

3103 + reason = _FLOAT16_NOT_SUPPORTED_MODELS[model_type]

3104 + raise ValueError(f"The model type {model_type!r} "

3105 + f"does not support float16. Reason: {reason}")

3106 +

3107 + return True

3108 +

3109 +

3110 +def _find_dtype(

3111 + model_id: str,

3112 config: PretrainedConfig,

3082 - dtype: Union[str, torch.dtype],

3083 -) -> torch.dtype:

3113 + *,

3114 + revision: Optional[str],

3115 +):

3116 # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct

3117 # because config.torch_dtype can be None.

3118 config_dtype = getattr(config, "torch_dtype", None)

@@ -3092,75 +3124,111 @@ def _get_and_verify_dtype(

3124 if config_dtype is None and hasattr(config, "vision_config"):

3125 config_dtype = getattr(config.vision_config, "torch_dtype", None)

3126

3127 + # Try to read the dtype of the weights if they are in safetensors format

3128 + if config_dtype is None:

3129 + repo_mt = try_get_safetensors_metadata(model_id, revision=revision)

3130 +

3131 + if repo_mt and (files_mt := repo_mt.files_metadata):

3132 + param_dtypes: set[torch.dtype] = {

3133 + _SAFETENSORS_TO_TORCH_DTYPE[dtype_str]

3134 + for file_mt in files_mt.values()

3135 + for dtype_str in file_mt.parameter_count

3136 + if dtype_str in _SAFETENSORS_TO_TORCH_DTYPE

3137 + }

3138 +

3139 + if param_dtypes:

3140 + return common_broadcastable_dtype(param_dtypes)

3141 +

3142 if config_dtype is None:

3143 config_dtype = torch.float32

3144

3098 - if isinstance(dtype, str):

3099 - dtype = dtype.lower()

3100 - if dtype == "auto":

3101 - # Set default dtype from model config

3102 - if config_dtype == torch.float32:

3103 - # Following common practice, we use float16 for float32 models

3104 - torch_dtype = torch.float16

3105 - else:

3106 - torch_dtype = config_dtype

3145 + return config_dtype

3146

3108 - if config.model_type == "plamo2":

3109 - logger.warning(

3110 - "For PLaMo2, we cast models to bfloat16 instead of using "

3111 - "float16 by default. This is because float16 does not work."

3112 - )

3113 - torch_dtype = torch.bfloat16

3147

3115 - # Deal with torch dtype fallback for device compatibility.

3116 - from vllm.platforms import current_platform

3117 - if torch_dtype not in current_platform.supported_dtypes:

3118 - device_name = current_platform.get_device_name()

3148 +def _resolve_auto_dtype(

3149 + model_type: str,

3150 + config_dtype: torch.dtype,

3151 + *,

3152 + is_pooling_model: bool,

3153 +):

3154 + from vllm.platforms import current_platform

3155

3120 - if ((capability := current_platform.get_device_capability())

3121 - is None):

3122 - compute_str = ""

3123 - else:

3124 - version_str = capability.as_version_str()

3125 - compute_str = f" (with compute capability {version_str})"

3126 - fallback_dtype = current_platform.supported_dtypes[0]

3127 - logger.warning(

3128 - "Your %s device%s doesn't support %s. " \

3129 - "Falling back to %s for compatibility.",

3130 - device_name, compute_str, torch_dtype, fallback_dtype

3131 - )

3132 - torch_dtype = fallback_dtype

3156 + supported_dtypes = [

3157 + dtype for dtype in current_platform.supported_dtypes

3158 + if _is_valid_dtype(model_type, dtype)

3159 + ]

3160

3134 - if current_platform.is_hpu() and torch_dtype == torch.float16:

3135 - logger.warning(

3136 - "For HPU, we cast models to bfloat16 instead of "

3137 - "using float16 by default. Please specify `dtype` if you "

3138 - "want to use float16.")

3139 - torch_dtype = torch.bfloat16

3140 - elif dtype == "float16" and config.model_type == "plamo2":

3141 - logger.warning(

3142 - "For PLaMo2, using float16 is unstable and might cause "

3143 - "unexpected behavior. Please use bfloat16 or float32 instead.")

3144 - torch_dtype = torch.float16

3161 + if is_pooling_model and torch.float16 in supported_dtypes:

3162 + preferred_dtype = torch.float16

3163 + else:

3164 + preferred_dtype = supported_dtypes[0]

3165 +

3166 + # Downcast for float32 models

3167 + if config_dtype == torch.float32:

3168 + config_dtype = preferred_dtype

3169 +

3170 + if config_dtype in supported_dtypes:

3171 + return config_dtype

3172 +

3173 + # Ensure device compatibility

3174 + device_name = current_platform.get_device_name()

3175 + device_capability = current_platform.get_device_capability()

3176 +

3177 + if device_capability is None:

3178 + device_str = f"{device_name!r}"

3179 + else:

3180 + version_str = device_capability.as_version_str()

3181 + device_str = f"{device_name!r} (with compute capability {version_str})"

3182 +

3183 + logger.warning(

3184 + "Your device %s doesn't support %s. "

3185 + "Falling back to %s for compatibility.",

3186 + device_str,

3187 + config_dtype,

3188 + preferred_dtype,

3189 + )

3190 +

3191 + return preferred_dtype

3192 +

3193 +

3194 +def _get_and_verify_dtype(

3195 + model_id: str,

3196 + config: PretrainedConfig,

3197 + dtype: Union[str, torch.dtype],

3198 + *,

3199 + is_pooling_model: bool,

3200 + revision: Optional[str] = None,

3201 +) -> torch.dtype:

3202 + config_dtype = _find_dtype(model_id, config, revision=revision)

3203 + model_type = config.model_type

3204 +

3205 + if isinstance(dtype, str):

3206 + dtype = dtype.lower()

3207 + if dtype == "auto":

3208 + # Set default dtype from model config

3209 + torch_dtype = _resolve_auto_dtype(

3210 + model_type,

3211 + config_dtype,

3212 + is_pooling_model=is_pooling_model,

3213 + )

3214 else:

3215 if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:

3147 - raise ValueError(f"Unknown dtype: {dtype}")

3216 + raise ValueError(f"Unknown dtype: {dtype!r}")

3217 torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]

3218 elif isinstance(dtype, torch.dtype):

3219 torch_dtype = dtype

3220 else:

3221 raise ValueError(f"Unknown dtype: {dtype}")

3222

3154 - # Verify the dtype.

3223 + _check_valid_dtype(model_type, torch_dtype)

3224 +

3225 if torch_dtype != config_dtype:

3226 if torch_dtype == torch.float32:

3227 # Upcasting to float32 is allowed.

3228 logger.info("Upcasting %s to %s.", config_dtype, torch_dtype)

3159 - pass

3229 elif config_dtype == torch.float32:

3230 # Downcasting from float32 to float16 or bfloat16 is allowed.

3231 logger.info("Downcasting %s to %s.", config_dtype, torch_dtype)

3163 - pass

3232 else:

3233 # Casting between float16 and bfloat16 is allowed with a warning.

3234 logger.warning("Casting %s to %s.", config_dtype, torch_dtype)

vllm/platforms/cpu.py

+1 -1

@@ -28,7 +28,7 @@ class CpuPlatform(Platform):

28 dispatch_key: str = "CPU"

30 @property

31 - def supported_dtypes(self) -> list:

31 + def supported_dtypes(self) -> list[torch.dtype]:

32 if self.get_cpu_architecture() == CpuArchEnum.POWERPC:

33 return [torch.bfloat16, torch.float32]

34 elif sys.platform.startswith(

vllm/transformers_utils/config.py

+33 -7

@@ -4,12 +4,12 @@ import enum

4 import json

5 import os

6 import time

7 -from functools import cache

7 +from functools import cache, partial

8 from pathlib import Path

9 -from typing import Any, Callable, Literal, Optional, Union

9 +from typing import Any, Callable, Literal, Optional, TypeVar, Union

11 import huggingface_hub

12 -from huggingface_hub import hf_hub_download

12 +from huggingface_hub import get_safetensors_metadata, hf_hub_download

13 from huggingface_hub import list_repo_files as hf_list_repo_files

14 from huggingface_hub import try_to_load_from_cache

15 from huggingface_hub.utils import (EntryNotFoundError, HfHubHTTPError,

@@ -93,10 +93,15 @@ class ConfigFormat(str, enum.Enum):

93 MISTRAL = "mistral"

96 -def with_retry(func: Callable[[], Any],

97 - log_msg: str,

98 - max_retries: int = 2,

99 - retry_delay: int = 2):

96 +_R = TypeVar("_R")

97 +

98 +

99 +def with_retry(

100 + func: Callable[[], _R],

101 + log_msg: str,

102 + max_retries: int = 2,

103 + retry_delay: int = 2,

104 +) -> _R:

105 for attempt in range(max_retries):

106 try:

107 return func()

@@ -109,6 +114,8 @@ def with_retry(func: Callable[[], Any],

114 time.sleep(retry_delay)

115 retry_delay *= 2

116

117 + raise AssertionError("Should not be reached")

118 +

119

120 # @cache doesn't cache exceptions

121 @cache

@@ -840,3 +847,22 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):

847 return resolve_obj_by_qualname(function_name)()

848 else:

849 return nn.Sigmoid() if config.num_labels == 1 else nn.Identity()

850 +

851 +

852 +def try_get_safetensors_metadata(

853 + model: str,

854 + *,

855 + revision: Optional[str] = None,

856 +):

857 + get_safetensors_metadata_partial = partial(

858 + get_safetensors_metadata,

859 + model,

860 + revision=revision,

861 + token=os.getenv('HF_TOKEN', None),

862 + )

863 +

864 + try:

865 + return with_retry(get_safetensors_metadata_partial,

866 + "Error retrieving safetensors")

867 + except Exception:

868 + return None

vllm/utils.py

+49 -2

@@ -37,8 +37,8 @@ from argparse import (Action, ArgumentDefaultsHelpFormatter, ArgumentParser,

37 _ArgumentGroup)

38 from asyncio import FIRST_COMPLETED, AbstractEventLoop, Task

39 from collections import UserDict, defaultdict

40 -from collections.abc import (AsyncGenerator, Awaitable, Generator, Hashable,

41 - Iterable, Iterator, KeysView, Mapping)

40 +from collections.abc import (AsyncGenerator, Awaitable, Collection, Generator,

41 + Hashable, Iterable, Iterator, KeysView, Mapping)

42 from concurrent.futures.process import ProcessPoolExecutor

43 from dataclasses import dataclass, field

44 from functools import cache, lru_cache, partial, wraps

@@ -979,6 +979,53 @@ def get_dtype_size(dtype: torch.dtype) -> int:

979 return torch.tensor([], dtype=dtype).element_size()

980

981

982 +# bool = 0, int = 1, float = 2, complex = 3

983 +def _get_precision_level(dtype: torch.dtype) -> int:

984 + # NOTE: Complex dtypes return `is_floating_point=False`

985 + return ((dtype != torch.bool) + dtype.is_floating_point +

986 + dtype.is_complex * 2)

987 +

988 +

989 +def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):

990 + """

991 + Test whether it is lossless to cast a tensor from

992 + `src_dtype` to `tgt_dtype`.

993 + """

994 + if src_dtype == tgt_dtype:

995 + return True

996 +

997 + src_level = _get_precision_level(src_dtype)

998 + tgt_level = _get_precision_level(tgt_dtype)

999 +

1000 + if src_level < tgt_level:

1001 + return True

1002 + if src_level > tgt_level:

1003 + return False

1004 +

1005 + # Compare integral types

1006 + if not src_dtype.is_floating_point and not src_dtype.is_complex:

1007 + src_info = torch.iinfo(src_dtype)

1008 + tgt_info = torch.iinfo(tgt_dtype)

1009 + return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max

1010 +

1011 + # Compare floating-point types

1012 + src_info = torch.finfo(src_dtype)

1013 + tgt_info = torch.finfo(tgt_dtype)

1014 + return (src_info.min >= tgt_info.min and src_info.max <= tgt_info.max

1015 + and src_info.resolution >= tgt_info.resolution)

1016 +

1017 +

1018 +def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):

1019 + """

1020 + Get the common `dtype` where all of the other `dtypes` can be

1021 + cast to it without losing any information.

1022 + """

1023 + return max(

1024 + dtypes,

1025 + key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes),

1026 + )

1027 +

1028 +

1029 # `collections` helpers

1030 def is_list_of(

1031 value: object,