xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 19:05:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

added support for quantize on LLM module (#1080)

Browse Source
This commit is contained in:
orellavie1212 2023-09-18 21:04:21 +03:00 committed by GitHub
parent 90979c38f8
commit fbe66e1d0b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/entrypoints/llm.py
View File

@ -38,6 +38,9 @@ class LLM:
However, if the `torch_dtype` in the config is `float32`, we will However, if the `torch_dtype` in the config is `float32`, we will
use `float16` instead. use `float16` instead.
seed: The seed to initialize the random number generator for sampling. seed: The seed to initialize the random number generator for sampling.
quantization: The method used to quantize the model weights. Currently,
we support "awq". If None, we assume the model weights are not
quantized and use `dtype` to determine the data type of the weights.
revision: The specific model version to use. It can be a branch name, revision: The specific model version to use. It can be a branch name,
a tag name, or a commit id. a tag name, or a commit id.
""" """
@ -51,6 +54,7 @@ class LLM:
tensor_parallel_size: int = 1, tensor_parallel_size: int = 1,
dtype: str = "auto", dtype: str = "auto",
seed: int = 0, seed: int = 0,
quantization: Optional[str] = None,
**kwargs, **kwargs,
) -> None: ) -> None:
if "disable_log_stats" not in kwargs: if "disable_log_stats" not in kwargs:
@ -63,6 +67,7 @@ class LLM:
tensor_parallel_size=tensor_parallel_size, tensor_parallel_size=tensor_parallel_size,
dtype=dtype, dtype=dtype,
seed=seed, seed=seed,
quantization=quantization,
**kwargs, **kwargs,
) )
self.llm_engine = LLMEngine.from_engine_args(engine_args) self.llm_engine = LLMEngine.from_engine_args(engine_args)