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, 02:15 UTC
[Bugfix] Fix handling of Tensorizer arguments for LoadConfig (#20643)
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Sanger Steel committed
Jul 9, 2025 at 11:36 UTC
4ac9c33f789d532e293cb2e14df0cd7814d2502d
4 files changed
+21
-52
tests/tensorizer_loader/test_tensorizer.py
-19
@@ -103,25 +103,6 @@ def write_keyfile(keyfile_path: str):
103
f.write(encryption_params.key)
104
105
106
-@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
107
-def test_can_deserialize_s3(vllm_runner):
108
- model_ref = "EleutherAI/pythia-1.4b"
109
- tensorized_path = f"s3://tensorized/{model_ref}/fp16/model.tensors"
110
-
111
- with vllm_runner(model_ref,
112
- load_format="tensorizer",
113
- model_loader_extra_config=TensorizerConfig(
114
- tensorizer_uri=tensorized_path,
115
- num_readers=1,
116
- s3_endpoint="object.ord1.coreweave.com",
117
- )) as loaded_hf_model:
118
- deserialized_outputs = loaded_hf_model.generate(
119
- prompts, sampling_params)
120
- # noqa: E501
121
-
122
- assert deserialized_outputs
123
-
124
-
106
@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
107
def test_deserialized_encrypted_vllm_model_has_same_outputs(
108
model_ref, vllm_runner, tmp_path, model_path):
vllm/engine/arg_utils.py
+15
-29
@@ -1003,41 +1003,27 @@ class EngineArgs:
1003
override_attention_dtype=self.override_attention_dtype,
1004
)
1005
1006
- def valid_tensorizer_config_provided(self) -> bool:
1007
- """
1008
- Checks if a parseable TensorizerConfig was passed to
1009
- self.model_loader_extra_config. It first checks if the config passed
1010
- is a dict or a TensorizerConfig object directly, and if the latter is
1011
- true (by checking that the object has TensorizerConfig's
1012
- .to_serializable() method), converts it in to a serializable dict
1013
- format
1014
- """
1015
- if self.model_loader_extra_config:
1016
- if hasattr(self.model_loader_extra_config, "to_serializable"):
1017
- self.model_loader_extra_config = (
1018
- self.model_loader_extra_config.to_serializable())
1019
- for allowed_to_pass in ["tensorizer_uri", "tensorizer_dir"]:
1020
- try:
1021
- self.model_loader_extra_config[allowed_to_pass]
1022
- return False
1023
- except KeyError:
1024
- pass
1025
- return True
1006
+ def validate_tensorizer_args(self):
1007
+ from vllm.model_executor.model_loader.tensorizer import (
1008
+ TensorizerConfig)
1009
+ for key in self.model_loader_extra_config:
1010
+ if key in TensorizerConfig._fields:
1011
+ self.model_loader_extra_config["tensorizer_config"][
1012
+ key] = self.model_loader_extra_config[key]
1013
1014
def create_load_config(self) -> LoadConfig:
1015
1016
if self.quantization == "bitsandbytes":
1017
self.load_format = "bitsandbytes"
1018
1032
- if (self.load_format == "tensorizer"
1033
- and self.valid_tensorizer_config_provided()):
1034
- logger.info("Inferring Tensorizer args from %s", self.model)
1035
- self.model_loader_extra_config = {"tensorizer_dir": self.model}
1036
- else:
1037
- logger.info(
1038
- "Using Tensorizer args from --model-loader-extra-config. "
1039
- "Note that you can now simply pass the S3 directory in the "
1040
- "model tag instead of providing the JSON string.")
1019
+ if self.load_format == "tensorizer":
1020
+ if hasattr(self.model_loader_extra_config, "to_serializable"):
1021
+ self.model_loader_extra_config = (
1022
+ self.model_loader_extra_config.to_serializable())
1023
+ self.model_loader_extra_config["tensorizer_config"] = {}
1024
+ self.model_loader_extra_config["tensorizer_config"][
1025
+ "tensorizer_dir"] = self.model
1026
+ self.validate_tensorizer_args()
1027
1028
return LoadConfig(
1029
load_format=self.load_format,
vllm/model_executor/model_loader/tensorizer.py
+5
-3
@@ -223,9 +223,11 @@ class TensorizerConfig(MutableMapping):
223
and re.search(r'%0\dd', self.tensorizer_uri) is not None
224
225
if self.tensorizer_dir and self.tensorizer_uri:
226
- raise ValueError(
227
- "Either tensorizer_dir or tensorizer_uri must be provided, "
228
- "not both.")
226
+ logger.warning_once(
227
+ "Provided both tensorizer_dir and tensorizer_uri. "
228
+ "Inferring tensorizer_dir from tensorizer_uri as the "
229
+ "latter takes precedence.")
230
+ self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
231
if self.tensorizer_dir and self.lora_dir:
232
raise ValueError(
233
"Only one of tensorizer_dir or lora_dir may be specified. "
vllm/model_executor/model_loader/tensorizer_loader.py
+1
-1
@@ -43,7 +43,7 @@ class TensorizerLoader(BaseModelLoader):
43
else:
44
validate_config(load_config.model_loader_extra_config)
45
self.tensorizer_config = TensorizerConfig(
46
- **load_config.model_loader_extra_config)
46
+ **load_config.model_loader_extra_config["tensorizer_config"])
47
48
def _verify_config(self, model_config: ModelConfig,
49
parallel_config: ParallelConfig):