From 8e5810fa2d88f38450022e90b370080b4d42389e Mon Sep 17 00:00:00 2001 From: ZX-ModelCloud Date: Fri, 31 Jul 2026 15:50:58 +0800 Subject: [PATCH] fix(runtime): support latest vLLM and SGLang APIs --- gptqmodel/models/loader.py | 84 +++++-- gptqmodel/utils/sglang.py | 383 +++++++++++++++++++++++++++++--- gptqmodel/utils/vllm.py | 385 ++++++++++++++++++++++++++------- tests/eval.py | 33 +-- tests/test_eval_loader_args.py | 36 ++- tests/test_vllm.py | 9 +- 6 files changed, 791 insertions(+), 139 deletions(-) diff --git a/gptqmodel/models/loader.py b/gptqmodel/models/loader.py index 41b7320fc..da3629eb0 100644 --- a/gptqmodel/models/loader.py +++ b/gptqmodel/models/loader.py @@ -88,6 +88,59 @@ ATTN_IMPLEMENTATION = "attn_implementation" +_EXTERNAL_BACKEND_FORMATS = { + BACKEND.VLLM: { + FORMAT.GPTQ, + FORMAT.GEMM, + }, + BACKEND.SGLANG: { + FORMAT.GPTQ, + FORMAT.GPTQ_V2, + FORMAT.GEMM, + FORMAT.MARLIN, + }, +} + + +def _validate_external_backend_format(backend: BACKEND, format_code: FORMAT) -> None: + supported_formats = _EXTERNAL_BACKEND_FORMATS.get(backend) + if supported_formats is None or format_code in supported_formats: + return + supported = ", ".join(f"FORMAT.{item.name}" for item in sorted(supported_formats, key=lambda item: item.name)) + raise ValueError(f"{backend} backend only supports {supported}: actual = {format_code}") + + +def _external_runtime_device_kwargs( + device: DEVICE, + requested_device_map: Optional[Union[str, Dict[str, Union[str, int]]]], +) -> Dict[str, Union[str, int]]: + runtime_kwargs: Dict[str, Union[str, int]] = {"device": device.type} + if not requested_device_map or device not in {DEVICE.CUDA, DEVICE.ROCM, DEVICE.XPU, DEVICE.NPU}: + return runtime_kwargs + + targets = ( + {requested_device_map} + if isinstance(requested_device_map, str) + else set(requested_device_map.values()) + ) + if len(targets) != 1: + return runtime_kwargs + target = targets.pop() + + if isinstance(target, int): + runtime_kwargs["base_gpu_id"] = target + return runtime_kwargs + if not isinstance(target, str) or target in {"cpu", "disk", "meta"}: + return runtime_kwargs + + try: + explicit_device = torch.device(target) + except (RuntimeError, ValueError): + return runtime_kwargs + if explicit_device.index is not None: + runtime_kwargs["base_gpu_id"] = explicit_device.index + return runtime_kwargs + def _should_print_module_tree() -> bool: """Keep expensive module-tree dumps opt-in during model loading.""" @@ -968,12 +1021,6 @@ def from_quantized( backend = normalize_backend(backend) device = auto_select_device(device, backend) - if backend == BACKEND.VLLM: - import os - - # to optimize vllm inference, set an environment variable 'VLLM_ATTENTION_BACKEND' to 'FLASHINFER'. - os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASHINFER' - model_local_path = get_model_local_path(model_id_or_path, **kwargs_without_internal) trust_remote_code = resolve_trust_remote_code(model_local_path, trust_remote_code=trust_remote_code) native_support = has_native_transformers_causallm_support(model_local_path) @@ -1116,15 +1163,15 @@ def from_quantized( if backend == BACKEND.VLLM or backend == BACKEND.SGLANG: runtime_generate = None - if backend == BACKEND.VLLM: - if format_code not in [FORMAT.GPTQ, FORMAT.GEMM]: - raise ValueError(f"{backend} backend only supports FORMAT.GPTQ or FORMAT.GEMM: actual = {qcfg.format}") - elif backend == BACKEND.SGLANG: - if format_code != FORMAT.GPTQ: - raise ValueError(f"{backend} backend only supports FORMAT.GPTQ: actual = {qcfg.format}") + _validate_external_backend_format(backend, format_code) if backend == BACKEND.VLLM: - from ..utils.vllm import load_model_by_vllm, vllm_generate + from ..utils.vllm import ( + get_vllm_device, + get_vllm_model_config, + load_model_by_vllm, + vllm_generate, + ) model = load_model_by_vllm( model=model_local_path, @@ -1132,18 +1179,23 @@ def from_quantized( **kwargs_without_internal, ) - model.config = model.llm_engine.model_config - model.device = model.llm_engine.vllm_config.device_config.device + model.config = get_vllm_model_config(model) + runtime_device = get_vllm_device(model) + if runtime_device is not None: + model.device = runtime_device runtime_generate = vllm_generate elif backend == BACKEND.SGLANG: from ..utils.sglang import load_model_by_sglang, sglang_generate + sglang_kwargs = dict(kwargs_without_internal) + for key, value in _external_runtime_device_kwargs(device, requested_device_map).items(): + sglang_kwargs.setdefault(key, value) model, hf_config = load_model_by_sglang( model=model_local_path, trust_remote_code=trust_remote_code, dtype=torch.float16, - **kwargs_without_internal, + **sglang_kwargs, ) model.config = hf_config runtime_generate = sglang_generate diff --git a/gptqmodel/utils/sglang.py b/gptqmodel/utils/sglang.py index c48074844..e99dc70a7 100644 --- a/gptqmodel/utils/sglang.py +++ b/gptqmodel/utils/sglang.py @@ -3,60 +3,389 @@ # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium -import multiprocessing as mp +from __future__ import annotations + +import inspect +from importlib.metadata import PackageNotFoundError, version +from typing import Any, Mapping, Optional import torch from transformers import AutoConfig +SGLANG_VERSION: Optional[str] = None +SGLANG_IMPORT_ERROR: Optional[Exception] = None +SGLANG_AVAILABLE = False +sgl = None + try: - import sglang as sgl + SGLANG_VERSION = version("sglang") +except PackageNotFoundError: + pass + +try: + import sglang as _sgl + + sgl = _sgl SGLANG_AVAILABLE = True -except ImportError: - SGLANG_AVAILABLE = False +except Exception as exc: + SGLANG_IMPORT_ERROR = exc + + +def _sglang_unavailable_message() -> str: + if SGLANG_VERSION is None and isinstance(SGLANG_IMPORT_ERROR, ModuleNotFoundError): + return "SGLang is not installed. Please install via `pip install -U 'sglang[srt]'`." + if SGLANG_IMPORT_ERROR is not None: + return ( + f"SGLang {SGLANG_VERSION or 'with unknown version'} is installed but failed to import: " + f"{type(SGLANG_IMPORT_ERROR).__name__}: {SGLANG_IMPORT_ERROR}." + ) + return "SGLang is not installed. Please install via `pip install -U 'sglang[srt]'`." + + +SGLANG_INSTALL_HINT = _sglang_unavailable_message() +_ENGINE_MARKER = "_gptqmodel_uses_sglang_engine" + + +def _require_sglang() -> None: + if not SGLANG_AVAILABLE: + raise ValueError(SGLANG_INSTALL_HINT) from SGLANG_IMPORT_ERROR + + +def _normalize_dtype(dtype: Any) -> Any: + if isinstance(dtype, torch.dtype): + return str(dtype).removeprefix("torch.") + if isinstance(dtype, str): + return dtype.removeprefix("torch.") + return dtype + + +def _move_alias(kwargs: dict[str, Any], source: str, target: str) -> None: + if source not in kwargs: + return + if target in kwargs: + raise ValueError(f"Pass only one of SGLang arguments `{source}` and `{target}`.") + kwargs[target] = kwargs.pop(source) + + +def _normalize_sglang_engine_kwargs(kwargs: Mapping[str, Any], trust_remote_code: bool) -> dict[str, Any]: + normalized = dict(kwargs) + _move_alias(normalized, "tensor_parallel_size", "tp_size") + _move_alias(normalized, "gpu_memory_utilization", "mem_fraction_static") + _move_alias(normalized, "max_model_len", "context_length") + _move_alias(normalized, "seed", "random_seed") + + if "enforce_eager" in normalized: + enforce_eager = normalized.pop("enforce_eager") + if "disable_cuda_graph" in normalized: + raise ValueError("Pass only one of SGLang arguments `enforce_eager` and `disable_cuda_graph`.") + normalized["disable_cuda_graph"] = bool(enforce_eager) + + dtype = normalized.get("dtype") + if dtype is not None: + normalized["dtype"] = _normalize_dtype(dtype) + + device = normalized.get("device") + if isinstance(device, torch.device): + if device.index is not None: + normalized.setdefault("base_gpu_id", device.index) + normalized["device"] = device.type + elif isinstance(device, str) and ":" in device: + parsed_device = torch.device(device) + if parsed_device.index is not None: + normalized.setdefault("base_gpu_id", parsed_device.index) + normalized["device"] = parsed_device.type + elif device == "rocm": + normalized["device"] = "cuda" + + for key in ("base_gpu_id", "context_length", "random_seed", "tp_size"): + if normalized.get(key) is not None: + normalized[key] = int(normalized[key]) + if normalized.get("mem_fraction_static") is not None: + normalized["mem_fraction_static"] = float(normalized["mem_fraction_static"]) + + normalized.setdefault("trust_remote_code", trust_remote_code) + return normalized -SGLANG_INSTALL_HINT = "sglang not installed. Please install via `pip install -U sglang`." def load_model_by_sglang( model, trust_remote_code, - **kwargs + **kwargs, ): - if not SGLANG_AVAILABLE: - raise ValueError(SGLANG_INSTALL_HINT) + _require_sglang() - mp.set_start_method('spawn') - runtime = sgl.Runtime( - model_path=model, - **kwargs, - ) - sgl.set_default_backend(runtime) hf_config = AutoConfig.from_pretrained( - model, trust_remote_code=trust_remote_code + model, + trust_remote_code=trust_remote_code, ) + runtime_kwargs = _normalize_sglang_engine_kwargs(kwargs, trust_remote_code) + engine_factory = getattr(sgl, "Engine", None) + if engine_factory is not None: + runtime = engine_factory( + model_path=model, + **runtime_kwargs, + ) + setattr(runtime, _ENGINE_MARKER, True) + else: + runtime = sgl.Runtime( + model_path=model, + **runtime_kwargs, + ) + setattr(runtime, _ENGINE_MARKER, False) + sgl.set_default_backend(runtime) return runtime, hf_config + if SGLANG_AVAILABLE: + @sgl.function - def generate(s, prompt, **kwargs): + def _legacy_generate(s, prompt, **kwargs): s += prompt s += sgl.gen("result", **kwargs) + else: - def generate(s, prompt, **kwargs): - print(SGLANG_INSTALL_HINT) -@torch.inference_mode + def _legacy_generate(s, prompt, **kwargs): + raise ValueError(SGLANG_INSTALL_HINT) + + +def _normalize_eos_token_ids(value: Any) -> list[int]: + if torch.is_tensor(value): + value = value.detach().cpu().tolist() + if isinstance(value, int) and not isinstance(value, bool): + return [value] + if isinstance(value, (list, tuple)) and all( + isinstance(item, int) and not isinstance(item, bool) for item in value + ): + return list(value) + raise TypeError("`eos_token_id` must be an integer or a sequence of integers.") + + +def _build_sglang_sampling_params(value: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + if value is None: + sampling_params = {} + elif isinstance(value, Mapping): + sampling_params = dict(value) + else: + raise TypeError("SGLang `sampling_params` must be a mapping.") + + if kwargs.get("max_length") is not None: + raise ValueError("SGLang does not support argument `max_length`. Please use `max_new_tokens` instead.") + if kwargs.get("min_length") is not None: + raise ValueError("SGLang does not support argument `min_length`. Please use `min_new_tokens` instead.") + + field_map = { + "num_return_sequences": "n", + "repetition_penalty": "repetition_penalty", + "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "min_p": "min_p", + "max_new_tokens": "max_new_tokens", + "max_tokens": "max_new_tokens", + "min_new_tokens": "min_new_tokens", + "min_tokens": "min_new_tokens", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "ignore_eos": "ignore_eos", + "stop": "stop", + "stop_token_ids": "stop_token_ids", + "regex": "regex", + "json_schema": "json_schema", + "sampling_seed": "sampling_seed", + } + for source, target in field_map.items(): + if kwargs.get(source) is not None: + sampling_params[target] = kwargs[source] + + if kwargs.get("eos_token_id") is not None: + eos_token_ids = _normalize_eos_token_ids(kwargs["eos_token_id"]) + existing_stop_ids = sampling_params.get("stop_token_ids") or [] + sampling_params["stop_token_ids"] = list(dict.fromkeys([*existing_stop_ids, *eos_token_ids])) + + if kwargs.get("do_sample") is False and "temperature" not in sampling_params: + sampling_params["temperature"] = 0.0 + return sampling_params + + +def _coerce_token_batch(value: Any) -> tuple[list[list[int]], bool]: + single_prompt = False + if torch.is_tensor(value): + single_prompt = value.ndim == 1 + value = value.detach().cpu().tolist() + elif isinstance(value, tuple): + value = list(value) + + if not isinstance(value, list) or not value: + raise ValueError("Token prompts must be a non-empty tensor or list.") + if all(isinstance(token, int) and not isinstance(token, bool) for token in value): + return [list(value)], True + + rows = [] + for row in value: + if torch.is_tensor(row): + row = row.detach().cpu().tolist() + elif isinstance(row, tuple): + row = list(row) + if not isinstance(row, list) or not row: + raise ValueError("Each token prompt must be a non-empty list of token IDs.") + if not all(isinstance(token, int) and not isinstance(token, bool) for token in row): + raise TypeError("Token prompts may only contain integer token IDs.") + rows.append(list(row)) + return rows, single_prompt + + +def _apply_attention_mask(token_batch: list[list[int]], attention_mask: Any) -> list[list[int]]: + if attention_mask is None: + return token_batch + if torch.is_tensor(attention_mask): + attention_mask = attention_mask.detach().cpu().tolist() + elif isinstance(attention_mask, tuple): + attention_mask = list(attention_mask) + if len(token_batch) == 1 and isinstance(attention_mask, list) and attention_mask: + if all(not isinstance(item, (list, tuple)) for item in attention_mask): + attention_mask = [attention_mask] + if not isinstance(attention_mask, list) or len(attention_mask) != len(token_batch): + raise ValueError("`attention_mask` must have the same batch size as `input_ids`.") + + filtered_batch = [] + for token_ids, mask in zip(token_batch, attention_mask): + if torch.is_tensor(mask): + mask = mask.detach().cpu().tolist() + elif isinstance(mask, tuple): + mask = list(mask) + if not isinstance(mask, list) or len(mask) != len(token_ids): + raise ValueError("Each `attention_mask` row must have the same length as its token prompt.") + filtered = [token_id for token_id, keep in zip(token_ids, mask) if bool(keep)] + if not filtered: + raise ValueError("`attention_mask` removed every token from a prompt.") + filtered_batch.append(filtered) + return filtered_batch + + +def _normalize_sglang_inputs(prompts: Any, input_ids: Any, attention_mask: Any): + if prompts is not None and input_ids is not None: + raise ValueError("Pass only one of `prompts` or `input_ids`.") + value = prompts if prompts is not None else input_ids + if value is None: + raise ValueError("Either prompts or input_ids must be provided.") + + if isinstance(value, str): + if input_ids is not None: + raise TypeError("`input_ids` cannot be a string.") + return value, None + if isinstance(value, (list, tuple)) and value and all(isinstance(item, str) for item in value): + if input_ids is not None: + raise TypeError("`input_ids` cannot contain strings.") + return list(value), None + + token_batch, single_prompt = _coerce_token_batch(value) + token_batch = _apply_attention_mask(token_batch, attention_mask) + return None, token_batch[0] if single_prompt else token_batch + + +def _extract_sglang_text(result: Any): + if isinstance(result, Mapping): + if "text" not in result: + raise RuntimeError("SGLang generation result is missing the `text` field.") + return result["text"] + if isinstance(result, list): + return [_extract_sglang_text(item) for item in result] + raise TypeError(f"Unexpected SGLang generation result type: {type(result)}.") + + +def _uses_engine_api(model: Any) -> bool: + marker = getattr(model, _ENGINE_MARKER, None) + if marker is not None: + return bool(marker) + return model.__class__.__module__.startswith("sglang.srt.entrypoints.") + + +def _legacy_sglang_sampling_params(sampling_params: Mapping[str, Any]) -> dict[str, Any]: + normalized = dict(sampling_params) + _move_alias(normalized, "max_new_tokens", "max_tokens") + _move_alias(normalized, "min_new_tokens", "min_tokens") + + try: + supported = set(inspect.signature(sgl.gen).parameters) + except (TypeError, ValueError): + supported = { + "frequency_penalty", + "ignore_eos", + "json_schema", + "max_tokens", + "min_p", + "min_tokens", + "n", + "presence_penalty", + "regex", + "stop", + "stop_token_ids", + "temperature", + "top_k", + "top_p", + } + unsupported = sorted(set(normalized) - supported) + if unsupported: + names = ", ".join(unsupported) + raise ValueError(f"The legacy SGLang Runtime frontend does not support sampling parameters: {names}.") + return normalized + + +@torch.inference_mode() def sglang_generate( - model, - **kwargs, + model, + **kwargs, ): - if not SGLANG_AVAILABLE: - raise ValueError(SGLANG_INSTALL_HINT) + _require_sglang() prompts = kwargs.pop("prompts", None) - state = generate.run( - prompt=prompts, - **kwargs, - ) + input_ids = kwargs.pop("input_ids", None) + attention_mask = kwargs.pop("attention_mask", None) + text_prompts, token_prompts = _normalize_sglang_inputs(prompts, input_ids, attention_mask) + + sampling_keys = { + "do_sample", + "eos_token_id", + "frequency_penalty", + "ignore_eos", + "json_schema", + "max_length", + "max_new_tokens", + "max_tokens", + "min_length", + "min_new_tokens", + "min_p", + "min_tokens", + "num_return_sequences", + "presence_penalty", + "regex", + "repetition_penalty", + "sampling_seed", + "stop", + "stop_token_ids", + "temperature", + "top_k", + "top_p", + } + sampling_params = _build_sglang_sampling_params(kwargs.pop("sampling_params", None), kwargs) + request_kwargs = {key: value for key, value in kwargs.items() if key not in sampling_keys} + request_kwargs.pop("pad_token_id", None) + if _uses_engine_api(model): + result = model.generate( + prompt=text_prompts, + input_ids=token_prompts, + sampling_params=sampling_params, + **request_kwargs, + ) + return _extract_sglang_text(result) + + if token_prompts is not None: + raise ValueError("The legacy SGLang Runtime frontend does not support `input_ids`; pass text prompts instead.") + state = _legacy_generate.run( + prompt=text_prompts, + **_legacy_sglang_sampling_params(sampling_params), + ) return state["result"] diff --git a/gptqmodel/utils/vllm.py b/gptqmodel/utils/vllm.py index e7f6c49a2..07cd2db4b 100644 --- a/gptqmodel/utils/vllm.py +++ b/gptqmodel/utils/vllm.py @@ -3,117 +3,348 @@ # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium -from typing import Any, Dict +from __future__ import annotations + +from importlib.metadata import PackageNotFoundError, version +from typing import Any, Dict, Mapping, Optional import torch +VLLM_VERSION: Optional[str] = None +VLLM_IMPORT_ERROR: Optional[Exception] = None +VLLM_AVAILABLE = False + +LLM = None +SamplingParams = None +TokensPrompt = None + try: - from vllm import LLM, SamplingParams, TokensPrompt + VLLM_VERSION = version("vllm") +except PackageNotFoundError: + pass +try: + from vllm import LLM as _LLM + from vllm import SamplingParams as _SamplingParams + + try: + from vllm import TokensPrompt as _TokensPrompt + except ImportError: + try: + from vllm.inputs import TokensPrompt as _TokensPrompt + except ImportError: + _TokensPrompt = None + + LLM = _LLM + SamplingParams = _SamplingParams + TokensPrompt = _TokensPrompt VLLM_AVAILABLE = True -except ImportError: - VLLM_AVAILABLE = False +except Exception as exc: + VLLM_IMPORT_ERROR = exc -VLLM_INSTALL_HINT = "vLLM not installed. Please install via `pip install -U vllm`." -# returns SamplingParams but we can't use this typehint since vLLM is optional depend -def convert_hf_params_to_vllm(hf_params: Dict[str, Any]): - if not VLLM_AVAILABLE: - raise ValueError(VLLM_INSTALL_HINT) - sampling_params = SamplingParams() +def _vllm_unavailable_message() -> str: + if VLLM_VERSION is None and isinstance(VLLM_IMPORT_ERROR, ModuleNotFoundError): + return "vLLM is not installed. Please install via `pip install -U vllm`." + if VLLM_IMPORT_ERROR is not None: + return ( + f"vLLM {VLLM_VERSION or 'with unknown version'} is installed but failed to import: " + f"{type(VLLM_IMPORT_ERROR).__name__}: {VLLM_IMPORT_ERROR}. " + "Check that vLLM's Python dependencies and its PyTorch/CUDA runtime build are mutually compatible." + ) + return "vLLM is not installed. Please install via `pip install -U vllm`." - if hf_params.get('num_return_sequences', None): - sampling_params.n = hf_params.get('num_return_sequences') - if hf_params.get('repetition_penalty', None): - sampling_params.repetition_penalty = hf_params.get('repetition_penalty') +VLLM_INSTALL_HINT = _vllm_unavailable_message() - if hf_params.get('temperature', None): - sampling_params.temperature = hf_params.get('temperature') - if hf_params.get('top_k', None): - sampling_params.top_k = hf_params.get('top_k') +def _require_vllm() -> None: + if not VLLM_AVAILABLE: + raise ValueError(VLLM_INSTALL_HINT) from VLLM_IMPORT_ERROR - if hf_params.get('top_p', None): - sampling_params.top_p = hf_params.get('top_p') - if hf_params.get('max_length', None): - raise ValueError("vLLM does not support argument `max_length`. Please use `max_new_tokens` instead.") - if hf_params.get('min_length', None): - raise ValueError("vLLM does not support argument `min_length`. Please use `min_new_tokens` instead.") +def _normalize_eos_token_ids(value: Any) -> list[int]: + if torch.is_tensor(value): + value = value.detach().cpu().tolist() + if isinstance(value, int) and not isinstance(value, bool): + return [value] + if isinstance(value, (list, tuple)) and all( + isinstance(item, int) and not isinstance(item, bool) for item in value + ): + return list(value) + raise TypeError("`eos_token_id` must be an integer or a sequence of integers.") - if hf_params.get('max_new_tokens', None): - sampling_params.max_tokens = hf_params.get('max_new_tokens') - if hf_params.get('min_new_tokens', None): - sampling_params.min_tokens = hf_params.get('min_new_tokens') +# Returns SamplingParams but we cannot use this type hint since vLLM is optional. +def convert_hf_params_to_vllm(hf_params: Dict[str, Any]): + _require_vllm() - if hf_params.get('eos_token_id', None): - sampling_params.stop_token_ids = [hf_params.get('eos_token_id'), None] + if hf_params.get("max_length") is not None: + raise ValueError("vLLM does not support argument `max_length`. Please use `max_new_tokens` instead.") + if hf_params.get("min_length") is not None: + raise ValueError("vLLM does not support argument `min_length`. Please use `min_new_tokens` instead.") + if hf_params.get("num_beams") not in (None, 1): + raise ValueError( + "GPT-QModel's vLLM generation adapter does not support Hugging Face beam search. " + "Use vLLM's beam-search API directly." + ) + if hf_params.get("do_sample") is False and (hf_params.get("num_return_sequences") or 1) > 1: + raise ValueError( + "vLLM requires sampling when `num_return_sequences` is greater than one. " + "Set `do_sample=True` or request a single sequence." + ) + + field_map = { + "num_return_sequences": "n", + "repetition_penalty": "repetition_penalty", + "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "min_p": "min_p", + "max_new_tokens": "max_tokens", + "max_tokens": "max_tokens", + "min_new_tokens": "min_tokens", + "min_tokens": "min_tokens", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "ignore_eos": "ignore_eos", + "stop": "stop", + "stop_token_ids": "stop_token_ids", + "seed": "seed", + } + sampling_kwargs = { + target: hf_params[source] + for source, target in field_map.items() + if hf_params.get(source) is not None + } + if hf_params.get("eos_token_id") is not None: + eos_token_ids = _normalize_eos_token_ids(hf_params["eos_token_id"]) + existing_stop_ids = sampling_kwargs.get("stop_token_ids") or [] + sampling_kwargs["stop_token_ids"] = list(dict.fromkeys([*existing_stop_ids, *eos_token_ids])) + if hf_params.get("do_sample") is False: + # Hugging Face ignores temperature for greedy decoding. vLLM selects + # greedy decoding by setting temperature to zero. + sampling_kwargs["temperature"] = 0.0 - return sampling_params + try: + return SamplingParams(**sampling_kwargs) + except TypeError as exc: + requested = ", ".join(sorted(sampling_kwargs)) + raise ValueError( + f"vLLM {VLLM_VERSION or 'unknown'} does not support the requested sampling parameters: {requested}." + ) from exc def load_model_by_vllm( - model, - **kwargs, + model, + **kwargs, ): - if not VLLM_AVAILABLE: - raise ValueError(VLLM_INSTALL_HINT) - - model = LLM( + _require_vllm() + return LLM( model=model, **kwargs, ) - return model + +def get_vllm_model_config(model): + model_config = getattr(model, "model_config", None) + if model_config is not None: + return model_config + + engine = getattr(model, "llm_engine", None) + if engine is None: + raise AttributeError("vLLM LLM instance is missing `llm_engine`.") + + vllm_config = getattr(engine, "vllm_config", None) or getattr(model, "vllm_config", None) + model_config = getattr(vllm_config, "model_config", None) + if model_config is None: + model_config = getattr(engine, "model_config", None) + if model_config is None: + raise AttributeError("vLLM engine exposes neither `vllm_config.model_config` nor `model_config`.") + return model_config + + +def get_vllm_device(model): + engine = getattr(model, "llm_engine", None) + if engine is None: + return None + + vllm_config = getattr(engine, "vllm_config", None) or getattr(model, "vllm_config", None) + device_config = getattr(vllm_config, "device_config", None) + if device_config is None: + device_config = getattr(engine, "device_config", None) + return getattr(device_config, "device", None) + + +def _coerce_token_batch(value: Any) -> list[list[int]]: + if torch.is_tensor(value): + value = value.detach().cpu().tolist() + elif isinstance(value, tuple): + value = list(value) + + if not isinstance(value, list) or not value: + raise ValueError("Token prompts must be a non-empty tensor or list.") + + if all(isinstance(token, int) and not isinstance(token, bool) for token in value): + return [list(value)] + + rows = [] + for row in value: + if torch.is_tensor(row): + row = row.detach().cpu().tolist() + elif isinstance(row, tuple): + row = list(row) + if not isinstance(row, list) or not row: + raise ValueError("Each token prompt must be a non-empty list of token IDs.") + if not all(isinstance(token, int) and not isinstance(token, bool) for token in row): + raise TypeError("Token prompts may only contain integer token IDs.") + rows.append(list(row)) + return rows + + +def _apply_attention_mask(token_batch: list[list[int]], attention_mask: Any) -> list[list[int]]: + if attention_mask is None: + return token_batch + if torch.is_tensor(attention_mask): + attention_mask = attention_mask.detach().cpu().tolist() + elif isinstance(attention_mask, tuple): + attention_mask = list(attention_mask) + + if len(token_batch) == 1 and isinstance(attention_mask, list) and attention_mask: + if all(not isinstance(item, (list, tuple)) for item in attention_mask): + attention_mask = [attention_mask] + + if not isinstance(attention_mask, list) or len(attention_mask) != len(token_batch): + raise ValueError("`attention_mask` must have the same batch size as `input_ids`.") + + filtered_batch = [] + for token_ids, mask in zip(token_batch, attention_mask): + if torch.is_tensor(mask): + mask = mask.detach().cpu().tolist() + elif isinstance(mask, tuple): + mask = list(mask) + if not isinstance(mask, list) or len(mask) != len(token_ids): + raise ValueError("Each `attention_mask` row must have the same length as its token prompt.") + filtered = [token_id for token_id, keep in zip(token_ids, mask) if bool(keep)] + if not filtered: + raise ValueError("`attention_mask` removed every token from a prompt.") + filtered_batch.append(filtered) + return filtered_batch + + +def _normalize_vllm_inputs(prompts: Any, input_ids: Any, attention_mask: Any): + if prompts is not None and input_ids is not None: + raise ValueError("Pass only one of `prompts` or `input_ids`.") + value = prompts if prompts is not None else input_ids + if value is None: + raise ValueError("Either prompts or input_ids must be provided.") + + if isinstance(value, str): + if input_ids is not None: + raise TypeError("`input_ids` cannot be a string.") + return value, None + if isinstance(value, (list, tuple)) and value and all(isinstance(item, str) for item in value): + if input_ids is not None: + raise TypeError("`input_ids` cannot contain strings.") + return list(value), None + + token_batch = _coerce_token_batch(value) + return None, _apply_attention_mask(token_batch, attention_mask) + + +def _build_sampling_params(value: Any, kwargs: Mapping[str, Any]): + if value is None: + hf_keys = ( + "num_return_sequences", + "repetition_penalty", + "temperature", + "top_k", + "top_p", + "min_p", + "max_length", + "min_length", + "max_new_tokens", + "max_tokens", + "min_new_tokens", + "min_tokens", + "eos_token_id", + "frequency_penalty", + "presence_penalty", + "ignore_eos", + "stop", + "stop_token_ids", + "seed", + "do_sample", + "num_beams", + ) + hf_params = {key: kwargs[key] for key in hf_keys if kwargs.get(key) is not None} + # GPT-QModel's public generate method follows Hugging Face semantics, + # where sampling is disabled unless explicitly requested. + hf_params.setdefault("do_sample", False) + return convert_hf_params_to_vllm(hf_params) + if isinstance(value, SamplingParams): + return value + if isinstance(value, Mapping): + try: + return SamplingParams(**dict(value)) + except TypeError as exc: + raise ValueError("Invalid vLLM `sampling_params` mapping.") from exc + raise TypeError("`sampling_params` must be a vLLM SamplingParams instance or a mapping.") + + +def _run_vllm_generation(model, text_prompts, token_batch, sampling_params, generation_kwargs): + if token_batch is None: + return model.generate(prompts=text_prompts, sampling_params=sampling_params, **generation_kwargs) + if TokensPrompt is not None: + token_prompts = [TokensPrompt(prompt_token_ids=prompt) for prompt in token_batch] + return model.generate(prompts=token_prompts, sampling_params=sampling_params, **generation_kwargs) + return model.generate(prompt_token_ids=token_batch, sampling_params=sampling_params, **generation_kwargs) @torch.inference_mode() def vllm_generate(model, **kwargs): - if not VLLM_AVAILABLE: - raise ValueError(VLLM_INSTALL_HINT) - - # Extract and validate prompts - prompts = kwargs.pop("prompts", None) or kwargs.pop("input_ids", None) - if prompts is None: - raise ValueError("Either prompts or input_ids must be provided") - - sampling_params = kwargs.get("sampling_params") - if not isinstance(sampling_params, SamplingParams): - hf_params = { - key: kwargs.get(key) for key in [ - 'num_return_sequences', 'repetition_penalty', 'temperature', - 'top_k', 'top_p', 'max_length', 'min_length', 'max_new_tokens', 'min_new_tokens', 'eos_token_id' - ] - } - sampling_params = convert_hf_params_to_vllm({k: v for k, v in hf_params.items() if v is not None}) - - # Convert prompts to vLLM format - if isinstance(prompts, torch.Tensor): - tokens_prompts = [TokensPrompt(prompt_token_ids=prompt) for prompt in prompts.tolist()] - req_results = model.generate(prompts=tokens_prompts, sampling_params=sampling_params) - elif isinstance(prompts, list): - if isinstance(prompts[0], list) or isinstance(prompts[0], int): - tokens_prompts = [TokensPrompt(prompt_token_ids=prompt) for prompt in prompts] - req_results = model.generate(prompts=tokens_prompts, sampling_params=sampling_params) - else: - req_results = model.generate(prompts=prompts, sampling_params=sampling_params) - elif isinstance(prompts, str): - req_results = model.generate(prompts=prompts, sampling_params=sampling_params) - else: - raise ValueError(f"Invalid input type for vllm_generate, type is {type(prompts)}") + _require_vllm() + + prompts = kwargs.pop("prompts", None) + input_ids = kwargs.pop("input_ids", None) + attention_mask = kwargs.pop("attention_mask", None) + text_prompts, token_batch = _normalize_vllm_inputs(prompts, input_ids, attention_mask) + + sampling_params = _build_sampling_params(kwargs.pop("sampling_params", None), kwargs) + generate_keys = ( + "use_tqdm", + "lora_request", + "priority", + "tokenization_kwargs", + "mm_processor_kwargs", + ) + generation_kwargs = {key: kwargs.pop(key) for key in generate_keys if key in kwargs} + req_results = _run_vllm_generation( + model, + text_prompts, + token_batch, + sampling_params, + generation_kwargs, + ) outputs = [] for result in req_results: - combined_token_ids = result.prompt_token_ids + list(result.outputs[0].token_ids) - outputs.append(combined_token_ids) + prompt_token_ids = list(result.prompt_token_ids) + for output in result.outputs: + outputs.append(prompt_token_ids + list(output.token_ids)) + if not outputs: + return torch.empty((0, 0), dtype=torch.long) - pad_token_id = model.get_tokenizer().pad_token_id + pad_token_id = kwargs.get("pad_token_id") + if pad_token_id is None: + tokenizer = model.get_tokenizer() + pad_token_id = tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = tokenizer.eos_token_id if pad_token_id is None: - pad_token_id = model.get_tokenizer().eos_token_id - max_length = max(len(sublist) for sublist in outputs) - padded_list = [sublist + [pad_token_id] * (max_length - len(sublist)) for sublist in outputs] + pad_token_id = 0 - return torch.Tensor(padded_list).to(torch.uint32) + max_length = max(len(output) for output in outputs) + padded = [output + [pad_token_id] * (max_length - len(output)) for output in outputs] + return torch.tensor(padded, dtype=torch.long) diff --git a/tests/eval.py b/tests/eval.py index d0b186e0f..5942d23f3 100644 --- a/tests/eval.py +++ b/tests/eval.py @@ -487,22 +487,23 @@ def _build_evalution_runtime( gpu_memory_utilization = engine_options.get("gpu_memory_utilization", 0.9) llm_kwargs = dict(engine_options.get("llm_kwargs", {}) or {}) - engine = evalution.VLLM( - dtype=engine_dtype, - batch_size=batch_size, - trust_remote_code=trust_remote_code, - padding_side=engine_padding_side, - seed=_evalution_engine_seed(engine_options), - tokenizer_mode=engine_options.get("tokenizer_mode", "auto"), - tensor_parallel_size=int(tensor_parallel_size), - gpu_memory_utilization=float(gpu_memory_utilization), - quantization=engine_options.get("quantization"), - max_model_len=int(max_model_len) if max_model_len is not None else None, - enforce_eager=bool(engine_options.get("enforce_eager", False)), - tokenizer_revision=engine_options.get("tokenizer_revision"), - vllm_path=engine_options.get("vllm_path"), - llm_kwargs=llm_kwargs, - ) + vllm_config = { + "dtype": engine_dtype, + "batch_size": batch_size, + "trust_remote_code": trust_remote_code, + "padding_side": engine_padding_side, + "seed": _evalution_engine_seed(engine_options), + "tokenizer_mode": engine_options.get("tokenizer_mode", "auto"), + "tensor_parallel_size": int(tensor_parallel_size), + "gpu_memory_utilization": float(gpu_memory_utilization), + "quantization": engine_options.get("quantization"), + "max_model_len": int(max_model_len) if max_model_len is not None else None, + "enforce_eager": bool(engine_options.get("enforce_eager", False)), + "tokenizer_revision": engine_options.get("tokenizer_revision"), + "vllm_path": engine_options.get("vllm_path"), + "llm_kwargs": llm_kwargs, + } + engine = safe_kwargs_call(evalution.VLLM, kwargs=vllm_config) else: sglang_config = _build_sglang_engine_kwargs( engine_options=engine_options, diff --git a/tests/test_eval_loader_args.py b/tests/test_eval_loader_args.py index 549cfc71f..e9f2cf3d8 100644 --- a/tests/test_eval_loader_args.py +++ b/tests/test_eval_loader_args.py @@ -68,8 +68,38 @@ def test_build_evalution_runtime_supports_vllm_engine_options(): captured = {} class FakeVLLM: - def __init__(self, **kwargs): - captured["engine_kwargs"] = kwargs + def __init__( + self, + *, + dtype=None, + batch_size=None, + trust_remote_code=None, + padding_side=None, + seed=None, + tokenizer_mode=None, + tensor_parallel_size=None, + gpu_memory_utilization=None, + quantization=None, + max_model_len=None, + enforce_eager=None, + tokenizer_revision=None, + llm_kwargs=None, + ): + captured["engine_kwargs"] = { + "dtype": dtype, + "batch_size": batch_size, + "trust_remote_code": trust_remote_code, + "padding_side": padding_side, + "seed": seed, + "tokenizer_mode": tokenizer_mode, + "tensor_parallel_size": tensor_parallel_size, + "gpu_memory_utilization": gpu_memory_utilization, + "quantization": quantization, + "max_model_len": max_model_len, + "enforce_eager": enforce_eager, + "tokenizer_revision": tokenizer_revision, + "llm_kwargs": llm_kwargs, + } def build(self, model_config): captured["model_config"] = model_config @@ -101,6 +131,7 @@ def to_dict(self): "quantization": "gptq", "tokenizer_mode": "auto", "max_model_len": "4096", + "vllm_path": "/tmp/legacy-vllm", "foo": "bar", }, tokenizer=None, @@ -115,6 +146,7 @@ def to_dict(self): assert captured["engine_kwargs"]["tensor_parallel_size"] == 2 assert captured["engine_kwargs"]["quantization"] == "gptq" assert captured["engine_kwargs"]["max_model_len"] == 4096 + assert "vllm_path" not in captured["engine_kwargs"] def test_build_evalution_runtime_supports_sglang_engine_options(): diff --git a/tests/test_vllm.py b/tests/test_vllm.py index 03a2eee5a..07d249eae 100644 --- a/tests/test_vllm.py +++ b/tests/test_vllm.py @@ -41,7 +41,14 @@ def setUpClass(self): ) try: - import vllm._C # noqa: F401,E402 + from vllm import LLM # noqa: F401,E402 + + try: + import vllm._C_stable_libtorch # noqa: F401,E402 + except ModuleNotFoundError: + # Compatibility with vLLM releases before the stable-libtorch + # extension became the default CUDA operator module. + import vllm._C # noqa: F401,E402 except Exception as exc: raise unittest.SkipTest(f"vllm runtime unavailable: {exc}") try: