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
from vllm.logger import init_logger
from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding)
from vllm.plugins import set_current_vllm_config
from vllm.utils import FlexibleArgumentParser
tensorizer_error_msg = None
......@@ -284,7 +285,8 @@ class TensorizerAgent:
model_args = self.tensorizer_config.hf_config
model_args.torch_dtype = self.tensorizer_config.dtype
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(
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