squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[bugfix] Fix example/tensorize_vllm_model tests (#10595)

Browse Source 
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
This commit is contained in:
Jee Jee Li 2024-11-24 09:22:33 +08:00 committed by GitHub
parent 04668ebe7a
commit 17d8fc1806
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/model_loader/tensorizer.py
View File

@ -19,6 +19,7 @@ from vllm.engine.llm_engine import LLMEngine
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding) VocabParallelEmbedding)
from vllm.plugins import set_current_vllm_config
from vllm.utils import FlexibleArgumentParser from vllm.utils import FlexibleArgumentParser
tensorizer_error_msg = None tensorizer_error_msg = None
@ -284,7 +285,8 @@ class TensorizerAgent:
model_args = self.tensorizer_config.hf_config model_args = self.tensorizer_config.hf_config
model_args.torch_dtype = self.tensorizer_config.dtype model_args.torch_dtype = self.tensorizer_config.dtype
assert self.tensorizer_config.model_class is not None assert self.tensorizer_config.model_class is not None
with no_init_or_tensor(): # TODO: Do we need to consider old-style model class?
with no_init_or_tensor(), set_current_vllm_config(self.vllm_config):
return self.tensorizer_config.model_class( return self.tensorizer_config.model_class(
vllm_config=self.vllm_config, ) vllm_config=self.vllm_config, )