Unverified Commit 17d8fc18 authored by Jee Jee Li's avatar Jee Jee Li Committed by GitHub
Browse files

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


Signed-off-by: default avatarJee Jee Li <pandaleefree@gmail.com>
parent 04668ebe
...@@ -19,6 +19,7 @@ from vllm.engine.llm_engine import LLMEngine ...@@ -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: ...@@ -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, )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment