mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
199 lines
8.6 KiB
Python
199 lines
8.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# Derived from vLLM; modified by DeepSeek.
|
||
|
||
from __future__ import annotations
|
||
|
||
import copy
|
||
import os
|
||
from collections import defaultdict
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
|
||
|
||
from deepinfer._C.engine import RequestOutput # pyright: ignore[reportMissingModuleSource]
|
||
from deepinfer.engine.arg_utils import EngineArgs
|
||
from deepinfer.engine.llm_engine import LLMEngine
|
||
from deepinfer.sampling_params import SamplingParams
|
||
from deepinfer.sequence import Sequence
|
||
from deepinfer.utils import Counter
|
||
|
||
|
||
class LLM:
|
||
"""An LLM for generating texts from given prompts and sampling parameters.
|
||
|
||
This class includes a tokenizer, a language model (possibly distributed
|
||
across multiple GPUs), and GPU memory space allocated for intermediate
|
||
states (aka KV cache). Given a batch of prompts and sampling parameters,
|
||
this class generates texts from the model, using an intelligent batching
|
||
mechanism and efficient memory management.
|
||
|
||
NOTE: For the comprehensive list of arguments, see `EngineArgs`.
|
||
|
||
Args:
|
||
model: The name or path of a HuggingFace Transformers model.
|
||
tokenizer: The name or path of a HuggingFace Transformers tokenizer.
|
||
trust_remote_code: Trust remote code (e.g., from HuggingFace) when
|
||
downloading the model and tokenizer.
|
||
revision: The specific model version to use. It can be a branch name,
|
||
a tag name, or a commit id.
|
||
tokenizer_revision: The specific tokenizer version to use. It can be a
|
||
branch name, a tag name, or a commit id.
|
||
seed: The seed to initialize the random number generator for sampling.
|
||
gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to
|
||
reserve for the model weights, activations, and KV cache. Higher
|
||
values will increase the KV cache size and thus improve the model's
|
||
throughput. However, if the value is too high, it may cause out-of-
|
||
memory (OOM) errors.
|
||
|
||
NOTE 开启 cudagraph 的时候，如果遇到'CUDA error: an illegal memory access was encountered ',
|
||
可能是 cudagraph OOM 了，可以调低 gpu_memory_utilization 再试试
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model: str,
|
||
tokenizer: str | None = None,
|
||
trust_remote_code: bool = False,
|
||
expert_parallel_size: int = 1,
|
||
revision: str | None = None,
|
||
tokenizer_revision: str | None = None,
|
||
seed: int = 0,
|
||
gpu_memory_utilization: float = 0.9,
|
||
decode_sequence: bool = True,
|
||
**kwargs: Any,
|
||
) -> None:
|
||
kwargs.setdefault("disable_log_stats", True)
|
||
kwargs.setdefault("stop_preempted_request", False)
|
||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||
engine_args = EngineArgs(
|
||
model=model,
|
||
tokenizer=tokenizer,
|
||
trust_remote_code=trust_remote_code,
|
||
expert_parallel_size=expert_parallel_size,
|
||
revision=revision,
|
||
tokenizer_revision=tokenizer_revision,
|
||
seed=seed,
|
||
gpu_memory_utilization=gpu_memory_utilization,
|
||
offline_inference=True,
|
||
**kwargs,
|
||
)
|
||
self.decode_sequence = decode_sequence
|
||
|
||
self.llm_engine = LLMEngine.from_engine_args(engine_args)
|
||
self.request_counter = Counter()
|
||
|
||
def get_tokenizer(self) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
|
||
return self.llm_engine.tokenizer
|
||
|
||
def generate(
|
||
self,
|
||
prompts: str | list[str] | None = None,
|
||
sampling_params: SamplingParams | list[SamplingParams] | None = None,
|
||
prompt_token_ids: list[list[int]] | list[np.ndarray[np.int32]] | None = None,
|
||
use_tqdm: bool = True,
|
||
multi_modal_data: dict | list[dict] | None = None,
|
||
) -> list[list[RequestOutput] | str]:
|
||
"""Generates the completions for the input prompts.
|
||
|
||
NOTE: This class automatically batches the given prompts, considering
|
||
the memory constraint. For the best performance, put all of your prompts
|
||
into a single list and pass it to this method.
|
||
|
||
Args:
|
||
prompts: A list of prompts to generate completions for.
|
||
sampling_params: The sampling parameters for text generation. If
|
||
None, we use the default sampling parameters.
|
||
prompt_token_ids: A list of token-ID lists (list[list[int]] | list[np.ndarray[np.int32]]) for the
|
||
prompts. If None, we use the tokenizer to convert the prompts to token IDs. Each item
|
||
is converted to np.int32 before being passed to the engine.
|
||
use_tqdm: Whether to use tqdm to display the progress bar.
|
||
|
||
Returns:
|
||
A list of `RequestOutput` objects containing the generated
|
||
completions in the same order as the input prompts.
|
||
"""
|
||
if prompts is None and prompt_token_ids is None:
|
||
raise ValueError("Either prompts or prompt_token_ids must be provided.")
|
||
if isinstance(prompts, str):
|
||
# Convert a single prompt to a list.
|
||
prompts = [prompts]
|
||
multi_modal_data = [multi_modal_data]
|
||
if prompts is not None and prompt_token_ids is not None and len(prompts) != len(prompt_token_ids):
|
||
raise ValueError("The lengths of prompts and prompt_token_ids must be the same.")
|
||
if sampling_params is None:
|
||
# Use default sampling params.
|
||
sampling_params = SamplingParams()
|
||
|
||
# Add requests to the engine.
|
||
num_requests = len(prompts) if prompts is not None else len(prompt_token_ids)
|
||
for i in range(num_requests):
|
||
prompt = prompts[i] if prompts is not None else None
|
||
token_ids = None if prompt_token_ids is None else prompt_token_ids[i]
|
||
|
||
self._add_request(
|
||
prompt,
|
||
# 共享同一个 SamplingParams 时无需在此复制：_add_request 内部会 deepcopy
|
||
sampling_params if isinstance(sampling_params, SamplingParams) else sampling_params[i],
|
||
token_ids,
|
||
None if multi_modal_data is None else multi_modal_data[i],
|
||
)
|
||
return self._run_engine(use_tqdm)
|
||
|
||
def _add_request(
|
||
self,
|
||
prompt: str | None,
|
||
sampling_params: SamplingParams,
|
||
prompt_token_ids: list[int] | np.ndarray[np.int32] | None = None,
|
||
multi_modal_data: dict | None = None,
|
||
) -> Sequence:
|
||
sampling_params = copy.deepcopy(sampling_params)
|
||
request_id = str(next(self.request_counter))
|
||
if prompt_token_ids is None:
|
||
if prompt is None:
|
||
raise ValueError("prompt and prompt_token_ids cannot both be None")
|
||
prompt_token_ids = self.llm_engine.tokenizer.encode(prompt)
|
||
prompt_token_ids = np.asarray(prompt_token_ids, dtype=np.int32)
|
||
if multi_modal_data is None and self.llm_engine.multi_modal_processor is not None:
|
||
multi_modal_data = dict(image_bytes=[], image_sizes=[])
|
||
return self.llm_engine.add_request(
|
||
user_id="0", # HACK 随便加一个 user_id
|
||
request_id=request_id,
|
||
sampling_params=sampling_params,
|
||
prompt_token_ids=prompt_token_ids,
|
||
multi_modal_data=multi_modal_data,
|
||
)
|
||
|
||
def _run_engine(self, use_tqdm: bool) -> list[list[RequestOutput] | str]:
|
||
# Initialize tqdm.
|
||
if use_tqdm:
|
||
num_requests = self.llm_engine.get_num_unfinished_requests()
|
||
pbar = tqdm(total=num_requests, desc="Processed prompts")
|
||
|
||
# Run the engine.
|
||
outputs: dict[str, list[RequestOutput]] = defaultdict(list)
|
||
while True:
|
||
unfinished = self.llm_engine.has_unfinished_requests()
|
||
step_outputs = self.llm_engine.step() if unfinished else self.llm_engine.flush_inflight_reqs()
|
||
for raw_output in step_outputs:
|
||
# clone step_outputs since step_outputs will be released in the next step
|
||
output = RequestOutput.clone(raw_output)
|
||
outputs[output.request_id].append(output)
|
||
if output.finished and use_tqdm:
|
||
pbar.update(1)
|
||
if not unfinished:
|
||
break
|
||
request_ids = sorted(outputs, key=int)
|
||
results: list[list[RequestOutput] | str] = [outputs[i] for i in request_ids]
|
||
if self.decode_sequence:
|
||
results = [self.llm_engine.tokenizer.decode([x.token_id for x in output]) for output in results]
|
||
|
||
if use_tqdm:
|
||
pbar.close()
|
||
|
||
return results
|
||
|
||
def close(self) -> None:
|
||
self.llm_engine.close()
|