llm.py 7.94 KB
Newer Older
1
from typing import List, Optional, Union
2
3

from tqdm import tqdm
Zhuohan Li's avatar
Zhuohan Li committed
4
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
5

Woosuk Kwon's avatar
Woosuk Kwon committed
6
7
8
9
10
from vllm.engine.arg_utils import EngineArgs
from vllm.engine.llm_engine import LLMEngine
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.utils import Counter
11
12
13


class LLM:
Woosuk Kwon's avatar
Woosuk Kwon committed
14
15
16
17
18
19
20
21
22
    """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: This class is intended to be used for offline inference. For online
23
    serving, use the `AsyncLLMEngine` class instead.
Zhuohan Li's avatar
Zhuohan Li committed
24
    NOTE: For the comprehensive list of arguments, see `EngineArgs`.
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
27

    Args:
        model: The name or path of a HuggingFace Transformers model.
28
        tokenizer: The name or path of a HuggingFace Transformers tokenizer.
29
30
        tokenizer_mode: The tokenizer mode. "auto" will use the fast tokenizer
            if available, and "slow" will always use the slow tokenizer.
31
32
        trust_remote_code: Trust remote code (e.g., from HuggingFace) when
            downloading the model and tokenizer.
Woosuk Kwon's avatar
Woosuk Kwon committed
33
34
35
        tensor_parallel_size: The number of GPUs to use for distributed
            execution with tensor parallelism.
        dtype: The data type for the model weights and activations. Currently,
Woosuk Kwon's avatar
Woosuk Kwon committed
36
37
38
39
            we support `float32`, `float16`, and `bfloat16`. If `auto`, we use
            the `torch_dtype` attribute specified in the model config file.
            However, if the `torch_dtype` in the config is `float32`, we will
            use `float16` instead.
40
41
42
        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.
Jasmond L's avatar
Jasmond L committed
43
44
        revision: The specific model version to use. It can be a branch name,
            a tag name, or a commit id.
45
46
47
48
49
50
51
52
53
54
55
        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.
        swap_space: The size (GiB) of CPU memory per GPU to use as swap space.
            This can be used for temporarily storing the states of the requests
            when their `best_of` sampling parameters are larger than 1. If all
            requests will have `best_of=1`, you can safely set this to 0.
            Otherwise, too small values may cause out-of-memory (OOM) errors.
Woosuk Kwon's avatar
Woosuk Kwon committed
56
    """
57
58
59
60

    def __init__(
        self,
        model: str,
61
        tokenizer: Optional[str] = None,
62
        tokenizer_mode: str = "auto",
63
        trust_remote_code: bool = False,
64
        tensor_parallel_size: int = 1,
Woosuk Kwon's avatar
Woosuk Kwon committed
65
        dtype: str = "auto",
66
        quantization: Optional[str] = None,
67
68
69
70
        revision: Optional[str] = None,
        seed: int = 0,
        gpu_memory_utilization: float = 0.9,
        swap_space: int = 4,
71
72
73
74
        **kwargs,
    ) -> None:
        if "disable_log_stats" not in kwargs:
            kwargs["disable_log_stats"] = True
Zhuohan Li's avatar
Zhuohan Li committed
75
        engine_args = EngineArgs(
76
            model=model,
77
            tokenizer=tokenizer,
78
            tokenizer_mode=tokenizer_mode,
79
            trust_remote_code=trust_remote_code,
80
81
            tensor_parallel_size=tensor_parallel_size,
            dtype=dtype,
82
            quantization=quantization,
83
84
85
86
            revision=revision,
            seed=seed,
            gpu_memory_utilization=gpu_memory_utilization,
            swap_space=swap_space,
87
88
            **kwargs,
        )
Zhuohan Li's avatar
Zhuohan Li committed
89
        self.llm_engine = LLMEngine.from_engine_args(engine_args)
90
91
        self.request_counter = Counter()

92
    def get_tokenizer(
93
            self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
Zhuohan Li's avatar
Zhuohan Li committed
94
        return self.llm_engine.tokenizer
95

96
97
98
99
100
101
    def set_tokenizer(
        self,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    ) -> None:
        self.llm_engine.tokenizer = tokenizer

102
103
    def generate(
        self,
Woosuk Kwon's avatar
Woosuk Kwon committed
104
        prompts: Optional[Union[str, List[str]]] = None,
105
        sampling_params: Optional[SamplingParams] = None,
106
        prompt_token_ids: Optional[List[List[int]]] = None,
107
108
        use_tqdm: bool = True,
    ) -> List[RequestOutput]:
Woosuk Kwon's avatar
Woosuk Kwon committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
        """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 IDs for the prompts. If None, we
                use the tokenizer to convert the prompts to token IDs.
            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.")
Woosuk Kwon's avatar
Woosuk Kwon committed
130
        if isinstance(prompts, str):
Woosuk Kwon's avatar
Woosuk Kwon committed
131
            # Convert a single prompt to a list.
Woosuk Kwon's avatar
Woosuk Kwon committed
132
            prompts = [prompts]
Woosuk Kwon's avatar
Woosuk Kwon committed
133
134
135
136
        if prompts is not None and prompt_token_ids is not None:
            if len(prompts) != len(prompt_token_ids):
                raise ValueError("The lengths of prompts and prompt_token_ids "
                                 "must be the same.")
137
        if sampling_params is None:
138
            # Use default sampling params.
139
            sampling_params = SamplingParams()
Woosuk Kwon's avatar
Woosuk Kwon committed
140

Zhuohan Li's avatar
Zhuohan Li committed
141
        # Add requests to the engine.
Woosuk Kwon's avatar
Woosuk Kwon committed
142
143
144
145
146
147
        if prompts is not None:
            num_requests = len(prompts)
        else:
            num_requests = len(prompt_token_ids)
        for i in range(num_requests):
            prompt = prompts[i] if prompts is not None else None
148
149
150
151
            if prompt_token_ids is None:
                token_ids = None
            else:
                token_ids = prompt_token_ids[i]
152
            self._add_request(prompt, sampling_params, token_ids)
Zhuohan Li's avatar
Zhuohan Li committed
153
        return self._run_engine(use_tqdm)
154

155
156
    def _add_request(
        self,
Woosuk Kwon's avatar
Woosuk Kwon committed
157
        prompt: Optional[str],
158
159
160
161
        sampling_params: SamplingParams,
        prompt_token_ids: Optional[List[int]],
    ) -> None:
        request_id = str(next(self.request_counter))
Zhuohan Li's avatar
Zhuohan Li committed
162
        self.llm_engine.add_request(request_id, prompt, sampling_params,
163
164
                                    prompt_token_ids)

Zhuohan Li's avatar
Zhuohan Li committed
165
    def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]:
166
167
        # Initialize tqdm.
        if use_tqdm:
Zhuohan Li's avatar
Zhuohan Li committed
168
            num_requests = self.llm_engine.get_num_unfinished_requests()
169
            pbar = tqdm(total=num_requests, desc="Processed prompts")
Zhuohan Li's avatar
Zhuohan Li committed
170
        # Run the engine.
171
        outputs: List[RequestOutput] = []
Zhuohan Li's avatar
Zhuohan Li committed
172
173
        while self.llm_engine.has_unfinished_requests():
            step_outputs = self.llm_engine.step()
174
            for output in step_outputs:
175
                if output.finished:
176
177
178
179
180
                    outputs.append(output)
                    if use_tqdm:
                        pbar.update(1)
        if use_tqdm:
            pbar.close()
181
182
183
184
        # Sort the outputs by request ID.
        # This is necessary because some requests may be finished earlier than
        # its previous requests.
        outputs = sorted(outputs, key=lambda x: int(x.request_id))
185
        return outputs