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 6, 17:00 UTC
Patch Mistral config (#37104)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
Julien Denize committed
Mar 16, 2026 at 13:22 UTC
ffbc2e5bdbfb7e4caae9c671696ca92fc9836101
3 files changed
+49
-30
vllm/transformers_utils/config.py
+38
-2
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
-from collections.abc import Callable
5
+from collections.abc import Callable, Iterator
6
+from contextlib import contextmanager
7
from dataclasses import asdict
8
from functools import cache, partial
9
from importlib.metadata import version
@@ -10,8 +11,10 @@ from pathlib import Path
11
from typing import Any, Literal, TypeAlias
12
13
import huggingface_hub
13
-from huggingface_hub import get_safetensors_metadata
14
+import torch
15
+from huggingface_hub import constants, get_safetensors_metadata
16
from packaging.version import Version
17
+from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
18
from transformers import GenerationConfig, PretrainedConfig
19
from transformers.models.auto.image_processing_auto import get_image_processor_config
20
from transformers.models.auto.modeling_auto import (
@@ -28,6 +31,7 @@ from vllm.transformers_utils.utils import (
31
parse_safetensors_file_metadata,
32
without_trust_remote_code,
33
)
34
+from vllm.utils.torch_utils import common_broadcastable_dtype
35
36
from .config_parser_base import ConfigParserBase
37
from .gguf_utils import (
@@ -135,6 +139,19 @@ def is_rope_parameters_nested(rope_parameters: dict[str, Any]) -> bool:
139
return set(rope_parameters.keys()).issubset(ALLOWED_ATTENTION_LAYER_TYPES)
140
141
142
+@contextmanager
143
+def _mistral_patch_hf_hub_constants() -> Iterator[None]:
144
+ hf_safetensors_single_file = constants.SAFETENSORS_SINGLE_FILE
145
+ hf_safetensors_index_file = constants.SAFETENSORS_INDEX_FILE
146
+ constants.SAFETENSORS_SINGLE_FILE = "consolidated.safetensors"
147
+ constants.SAFETENSORS_INDEX_FILE = "consolidated.safetensors.index.json"
148
+ try:
149
+ yield
150
+ finally:
151
+ constants.SAFETENSORS_SINGLE_FILE = hf_safetensors_single_file
152
+ constants.SAFETENSORS_INDEX_FILE = hf_safetensors_index_file
153
+
154
+
155
class HFConfigParser(ConfigParserBase):
156
def parse(
157
self,
@@ -245,6 +262,25 @@ class MistralConfigParser(ConfigParserBase):
262
except OSError: # Not found
263
hf_config_dict = {}
264
265
+ if config_dict.get("dtype") is None:
266
+ with _mistral_patch_hf_hub_constants():
267
+ model_str = model if isinstance(model, str) else model.as_posix()
268
+ param_mt = get_safetensors_params_metadata(model_str, revision=revision)
269
+ if param_mt:
270
+ param_dtypes: set[torch.dtype] = {
271
+ _SAFETENSORS_TO_TORCH_DTYPE[dtype]
272
+ for info in param_mt.values()
273
+ if (dtype := info.get("dtype", None))
274
+ and dtype in _SAFETENSORS_TO_TORCH_DTYPE
275
+ }
276
+
277
+ if param_dtypes:
278
+ config_dict["dtype"] = common_broadcastable_dtype(param_dtypes)
279
+ logger.info_once(
280
+ "Inferred from consolidated*.safetensors files "
281
+ f"{config_dict['dtype']} dtype."
282
+ )
283
+
284
config = adapt_config_dict(config_dict, defaults=hf_config_dict)
285
286
return config_dict, config
vllm/transformers_utils/configs/mistral.py
+10
-7
@@ -113,12 +113,13 @@ def _remap_mistral_vision_args(config: dict) -> dict:
113
114
def _remap_mistral_yarn_args(config: dict) -> dict:
115
yarn_config_map = {
116
- "factor": "factor",
117
- "original_max_position_embeddings": "original_max_position_embeddings",
118
- "beta": "beta_fast",
119
- "alpha": "beta_slow",
120
- "apply_scale": "apply_yarn_scaling",
116
+ "factor": ("factor", float),
117
+ "original_max_position_embeddings": ("original_max_position_embeddings", int),
118
+ "beta": ("beta_fast", float),
119
+ "alpha": ("beta_slow", float),
120
+ "apply_scale": ("apply_yarn_scaling", bool),
121
}
122
+
123
yarn_config = config.get("yarn") or {}
124
config["rope_parameters"] = {
125
"rope_type": "yarn",
@@ -128,9 +129,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
129
if rope_theta := config.pop("rope_theta", None):
130
config["rope_parameters"]["rope_theta"] = rope_theta
131
131
- for old_name, new_name in yarn_config_map.items():
132
+ for old_name, (new_name, cast) in yarn_config_map.items():
133
if old_name in yarn_config:
133
- config["rope_parameters"][new_name] = yarn_config.pop(old_name)
134
+ # Cast to remove Transformers > v5 type warnings
135
+ config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name))
136
137
assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
138
@@ -154,6 +156,7 @@ def _remap_general_mistral_args(config: dict) -> dict:
156
"tie_word_embeddings": ("tied_embeddings", False),
157
"max_seq_len": ("max_seq_len", config.get("max_position_embeddings", 128_000)),
158
"max_position_embeddings": ("max_position_embeddings", 128_000),
159
+ "dtype": ("dtype", config.get("dtype")),
160
}
161
162
for key, new_key in config_mapping.items():
vllm/transformers_utils/model_arch_config_convertor.py
+1
-21
@@ -1,12 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Iterator
5
-from contextlib import contextmanager
4
from typing import final
5
6
import torch
9
-from huggingface_hub import constants
7
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
8
from transformers import PretrainedConfig
9
@@ -25,22 +22,6 @@ from vllm.utils.torch_utils import common_broadcastable_dtype
22
logger = init_logger(__name__)
23
24
28
-@contextmanager
29
-def _maybe_patch_hf_hub_constants(config_format: ConfigFormat) -> Iterator[None]:
30
- if config_format == "mistral":
31
- hf_safetensors_single_file = constants.SAFETENSORS_SINGLE_FILE
32
- hf_safetensors_index_file = constants.SAFETENSORS_INDEX_FILE
33
- constants.SAFETENSORS_SINGLE_FILE = "consolidated.safetensors"
34
- constants.SAFETENSORS_INDEX_FILE = "consolidated.safetensors.index.json"
35
- try:
36
- yield
37
- finally:
38
- constants.SAFETENSORS_SINGLE_FILE = hf_safetensors_single_file
39
- constants.SAFETENSORS_INDEX_FILE = hf_safetensors_index_file
40
- else:
41
- yield
42
-
43
-
25
class ModelArchConfigConvertorBase:
26
def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):
27
self.hf_config = hf_config
@@ -164,8 +145,7 @@ class ModelArchConfigConvertorBase:
145
146
# Try to read the dtype of the weights if they are in safetensors format
147
if config_dtype is None:
167
- with _maybe_patch_hf_hub_constants(config_format):
168
- param_mt = get_safetensors_params_metadata(model_id, revision=revision)
148
+ param_mt = get_safetensors_params_metadata(model_id, revision=revision)
149
150
if param_mt:
151
param_dtypes: set[torch.dtype] = {