vllm_causallms.py 16.7 KB
Newer Older
1
2
3
4
5
6
import copy
from importlib.util import find_spec
from typing import List, Literal, Optional, Tuple, Union

from tqdm import tqdm

baberabb's avatar
baberabb committed
7
8
9
from lm_eval.api.instance import Instance
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
10
from lm_eval.models.utils import Collator, divide
11
12
13
14
15
from lm_eval.utils import (
    eval_logger,
    get_rolling_token_windows,
    make_disjoint_window,
)
16

Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
17

18
try:
19
    import ray
baberabb's avatar
baberabb committed
20
    from ray.util.multiprocessing import Pool
21
    from vllm import LLM, SamplingParams
baberabb's avatar
baberabb committed
22
    from vllm.transformers_utils.tokenizer import get_tokenizer
23
24
except ModuleNotFoundError:
    pass
Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
25

26
eval_logger = eval_logger
baberabb's avatar
baberabb committed
27

baberabb's avatar
baberabb committed
28

baberabb's avatar
baberabb committed
29
# adapted from https://github.com/vllm-project/vllm/issues/367#issuecomment-1788341727
30
31
32
def run_inference_one_model(
    model_args: dict, sampling_params, requests: List[List[int]]
):
baberabb's avatar
baberabb committed
33
34
35
36
    llm = LLM(**model_args)
    return llm.generate(prompt_token_ids=requests, sampling_params=sampling_params)


baberabb's avatar
baberabb committed
37
38
39
40
41
42
43
44
45
46
@register_model("vllm")
class VLLM(LM):
    _DEFAULT_MAX_LENGTH = 2048

    def __init__(
        self,
        pretrained="gpt2",
        dtype: Literal["float16", "bfloat16", "float32", "auto"] = "auto",
        revision: Optional[str] = None,
        trust_remote_code: Optional[bool] = False,
baberabb's avatar
baberabb committed
47
        tokenizer: Optional[str] = None,
baberabb's avatar
baberabb committed
48
        tokenizer_mode: Literal["auto", "slow"] = "auto",
baberabb's avatar
baberabb committed
49
        tokenizer_revision: Optional[str] = None,
baberabb's avatar
baberabb committed
50
        tensor_parallel_size: int = 1,
51
        quantization: Optional[str] = None,
baberabb's avatar
baberabb committed
52
53
        max_gen_toks: int = 256,
        swap_space: int = 4,
baberabb's avatar
baberabb committed
54
        batch_size: Union[str, int] = 1,
baberabb's avatar
baberabb committed
55
        max_batch_size=None,
baberabb's avatar
baberabb committed
56
        max_length: int = None,
57
        max_model_len: int = None,
baberabb's avatar
baberabb committed
58
        seed: int = 1234,
59
        gpu_memory_utilization: float = 0.9,
baberabb's avatar
baberabb committed
60
        device: str = "cuda",
61
        data_parallel_size: int = 1,
baberabb's avatar
baberabb committed
62
63
    ):
        super().__init__()
64

65
        if not find_spec("vllm"):
66
            raise Exception(
67
68
                "attempted to use 'vllm' LM type, but package `vllm` is not installed. "
                "Please install vllm via `pip install lm-eval[vllm]` or `pip install -e .[vllm]`"
Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
69
70
            )

baberabb's avatar
baberabb committed
71
        assert "cuda" in device or device is None, "vLLM only supports CUDA"
72
73
74
75
76
        assert (
            max_length is None or max_model_len is None
        ), "Either max_length or max_model_len may be provided, but not both"

        self._max_length = max_model_len if max_model_len is not None else max_length
baberabb's avatar
baberabb committed
77
        self.tensor_parallel_size = int(tensor_parallel_size)
78
        self.data_parallel_size = int(data_parallel_size)
baberabb's avatar
baberabb committed
79
80
81
82
83
        self.model_args = {
            "model": pretrained,
            "gpu_memory_utilization": float(gpu_memory_utilization),
            "revision": revision,
            "dtype": dtype,
baberabb's avatar
baberabb committed
84
            "tokenizer": tokenizer,
baberabb's avatar
baberabb committed
85
            "tokenizer_mode": tokenizer_mode,
baberabb's avatar
baberabb committed
86
            "tokenizer_revision": tokenizer_revision,
baberabb's avatar
baberabb committed
87
88
            "trust_remote_code": trust_remote_code,
            "tensor_parallel_size": int(tensor_parallel_size),
89
            "max_model_len": int(self._max_length) if self._max_length else None,
baberabb's avatar
baberabb committed
90
91
92
93
            "swap_space": int(swap_space),
            "quantization": quantization,
            "seed": int(seed),
        }
94
95
96
97
98
        self.batch_size = (
            "auto"
            if isinstance(batch_size, str) and "auto" in batch_size
            else batch_size
        )
99
        if self.data_parallel_size <= 1:
baberabb's avatar
baberabb committed
100
            self.model = LLM(**self.model_args)
baberabb's avatar
baberabb committed
101
102
        else:
            self.model_args["worker_use_ray"] = True
103
104
105
106
107
108
109
110
            self.batch_size = "auto"
            eval_logger.info("Manual batching is not compatible with data parallelism.")

            from transformers import AutoConfig

            self._config = AutoConfig.from_pretrained(
                pretrained, trust_remote_code=trust_remote_code, revision=revision
            )
baberabb's avatar
nits  
baberabb committed
111
112
113
114
115
116
        self.tokenizer = get_tokenizer(
            tokenizer if tokenizer else pretrained,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=trust_remote_code,
            tokenizer_revision=tokenizer_revision,
        )
117

baberabb's avatar
baberabb committed
118
119
120
121
122
123
124
125
126
127
128
        self._max_gen_toks = max_gen_toks

    @property
    def eot_token_id(self):
        # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
        return self.tokenizer.eos_token_id

    @property
    def max_length(self):
        if self._max_length:  # if max length manually set, return it
            return self._max_length
129
130
131
132
133
134
135
136
137
138
139
140
        if self.data_parallel_size <= 1:
            return self.model.llm_engine.model_config.max_model_len
        else:
            seqlen_config_attrs = ("n_positions", "max_position_embeddings", "n_ctx")
            for attr in seqlen_config_attrs:
                if hasattr(self._config, attr):
                    return getattr(self._config, attr)
            if hasattr(self.tokenizer, "model_max_length"):
                if self.tokenizer.model_max_length == 1000000000000000019884624838656:
                    return self._DEFAULT_MAX_LENGTH
                return self.tokenizer.model_max_length
            return self._DEFAULT_MAX_LENGTH
baberabb's avatar
baberabb committed
141
142
143
144
145

    @property
    def max_gen_toks(self):
        return self._max_gen_toks

baberabb's avatar
baberabb committed
146
147
148
149
150
151
152
    def tok_encode(
        self,
        string: str,
        left_truncate_len=None,
        add_special_tokens=False,
        truncation=False,
    ):
baberabb's avatar
baberabb committed
153
        """ """
baberabb's avatar
baberabb committed
154
155
156
        encoding = self.tokenizer.encode(
            string, add_special_tokens=add_special_tokens, truncation=truncation
        )
baberabb's avatar
baberabb committed
157
158
159
160
161
162
163
164
165

        # left-truncate the encoded context to be at most `left_truncate_len` tokens long
        if left_truncate_len:
            encoding = encoding[-left_truncate_len:]

        return encoding

    def _model_generate(
        self,
baberabb's avatar
baberabb committed
166
        requests: List[List[int]] = None,
baberabb's avatar
baberabb committed
167
168
169
170
171
172
        generate: bool = False,
        max_tokens: int = None,
        stop: Optional[List[str]] = None,
        **kwargs,
    ):
        if generate:
173
            kwargs = self.modify_gen_kwargs(kwargs)
baberabb's avatar
baberabb committed
174
            sampling_params = SamplingParams(max_tokens=max_tokens, stop=stop, **kwargs)
baberabb's avatar
baberabb committed
175
        else:
baberabb's avatar
baberabb committed
176
            sampling_params = SamplingParams(
177
                temperature=0, prompt_logprobs=1, max_tokens=1
baberabb's avatar
baberabb committed
178
            )
179
        if self.data_parallel_size > 1:
180
            requests = [list(x) for x in divide(requests, self.data_parallel_size)]
baberabb's avatar
baberabb committed
181
            inputs = [(self.model_args, sampling_params, req) for req in requests]
baberabb's avatar
baberabb committed
182

183
            with Pool(self.data_parallel_size) as pool:
baberabb's avatar
baberabb committed
184
                results = pool.starmap(run_inference_one_model, inputs)
185
186
            # Invoke ray.shutdown() to prevent hang-ups if subsequent calls required.
            ray.shutdown()
baberabb's avatar
baberabb committed
187
188
189
190
191
192
            # flatten results
            return [item for sublist in results for item in sublist]

        outputs = self.model.generate(
            prompt_token_ids=requests,
            sampling_params=sampling_params,
193
            use_tqdm=True if self.batch_size == "auto" else False,
baberabb's avatar
baberabb committed
194
        )
baberabb's avatar
baberabb committed
195
196
        return outputs

baberabb's avatar
baberabb committed
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
    def _encode_pair(
        self, context: str, continuation: str
    ) -> Tuple[List[int], List[int]]:
        n_spaces = len(context) - len(context.rstrip())
        if n_spaces > 0:
            continuation = context[-n_spaces:] + continuation
            context = context[:-n_spaces]

        whole_enc = self.tok_encode(context + continuation, add_special_tokens=False)
        context_enc = self.tok_encode(context, add_special_tokens=False)

        context_enc_len = len(context_enc)
        continuation_enc = whole_enc[context_enc_len:]
        return context_enc, continuation_enc

baberabb's avatar
baberabb committed
212
    def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
baberabb's avatar
baberabb committed
213
214
215
216
        new_reqs = []
        for context, continuation in [req.args for req in requests]:
            if context == "":
                # end of text as context
217
218
219
                context_enc, continuation_enc = (
                    [self.eot_token_id],
                    self.tok_encode(continuation),
baberabb's avatar
baberabb committed
220
                )
baberabb's avatar
baberabb committed
221
            else:
baberabb's avatar
baberabb committed
222
                context_enc, continuation_enc = self._encode_pair(context, continuation)
baberabb's avatar
baberabb committed
223
224
225
226
227

            new_reqs.append(((context, continuation), context_enc, continuation_enc))

        return self._loglikelihood_tokens(new_reqs)

baberabb's avatar
baberabb committed
228
    def loglikelihood_rolling(self, requests: List[Instance]) -> List[float]:
baberabb's avatar
baberabb committed
229
230
231
232
233
        loglikelihoods = []

        for (string,) in tqdm([req.args for req in requests]):
            rolling_token_windows = list(
                map(
234
235
                    make_disjoint_window,
                    get_rolling_token_windows(
baberabb's avatar
baberabb committed
236
237
                        token_list=self.tok_encode(string),
                        prefix_token=self.eot_token_id,
baberabb's avatar
baberabb committed
238
                        max_seq_len=self.max_length - 1,
baberabb's avatar
baberabb committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
                        context_len=1,
                    ),
                )
            )

            rolling_token_windows = [(None,) + x for x in rolling_token_windows]

            string_nll = self._loglikelihood_tokens(
                rolling_token_windows,
            )

            # discard is_greedy
            string_nll = [x[0] for x in string_nll]

            string_nll = sum(string_nll)
            loglikelihoods.append(string_nll)
        return loglikelihoods

    def generate_until(self, requests: List[Instance]) -> List[str]:
258
        res = []
baberabb's avatar
baberabb committed
259
260
261

        # batch tokenize contexts
        context, all_gen_kwargs = zip(*(req.args for req in requests))
262
        context_encoding = self.tokenizer(context, add_special_tokens=False).input_ids
baberabb's avatar
baberabb committed
263
264
265
        requests = [
            ((a, b), c) for a, b, c in zip(context, context_encoding, all_gen_kwargs)
        ]
baberabb's avatar
baberabb committed
266
267
268
269
270
271
272
273

        def _collate_gen(_requests):
            # the negative sign on len(toks) sorts descending - this has a few advantages:
            # - time estimates will always be over not underestimates, which is more useful for planning
            # - to know the size of a batch when going through the list, you know the first one is always the batch
            #   padded context length. this is useful to simplify the batching logic and more importantly to make
            #   automatic adaptive batches much much easier to implement
            # - any OOMs will happen right away rather than near the end
274
            return -len(_requests[0][1]), _requests[0][0]
baberabb's avatar
baberabb committed
275
276
277
278

        # we group requests by their generation_kwargs,
        # so that we don't try to execute e.g. greedy sampling and temp=0.8 sampling
        # in the same batch.
279
280
281
282
        re_ords = Collator(requests, _collate_gen, grouping=True)
        chunks = re_ords.get_batched(
            n=int(self.batch_size) if self.batch_size != "auto" else 0, batch_fn=None
        )
baberabb's avatar
baberabb committed
283
284
285

        pbar = tqdm(total=len(requests), disable=(self.rank != 0))
        # for each different set of kwargs, we execute all requests, by batch.
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        for chunk in chunks:
            context_and_encoding, all_gen_kwargs = zip(*chunk)
            context, context_encoding = zip(*context_and_encoding)
            # we assume all gen kwargs in the batch are the same
            # this is safe to assume because the `grouper` object ensures it.
            gen_kwargs = all_gen_kwargs[0]
            # unpack our keyword arguments.
            until = None
            if isinstance(gen_kwargs, dict):
                kwargs = copy.deepcopy(gen_kwargs)  # edge case for repeats > 1
                if "until" in kwargs.keys():
                    until = kwargs.pop("until")
                    if isinstance(until, str):
                        until = [until]
                    elif not isinstance(until, list):
                        raise ValueError(
                            f"Expected `kwargs['until']` to be of type Union[str,list] but got {until}"
                        )
            else:
                raise ValueError(
                    f"Expected `kwargs` to be of type `dict` but got {gen_kwargs}"
baberabb's avatar
baberabb committed
307
                )
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
            if not until:
                until = [self.tokenizer.decode(self.eot_token_id)]
            if "max_gen_toks" in kwargs.keys():
                max_gen_toks = kwargs.pop("max_gen_toks")
            else:
                max_gen_toks = self.max_gen_toks

            # set the max length in tokens of inputs ("context_enc")
            # max len for inputs = max length, minus room to generate the max new tokens
            max_ctx_len = self.max_length - max_gen_toks
            context_encoding = [x[-max_ctx_len:] for x in context_encoding]

            # perform batched generation
            cont = self._model_generate(
                requests=context_encoding,
                generate=True,
                max_tokens=max_gen_toks,
                stop=until,
                **kwargs,
            )
baberabb's avatar
baberabb committed
328

329
330
331
332
333
334
335
336
            # cache generations
            for output, context in zip(cont, context):
                generated_text = output.outputs[0].text
                res.append(generated_text)
                self.cache_hook.add_partial(
                    "generate_until", (context, gen_kwargs), generated_text
                )
                pbar.update(1)
baberabb's avatar
baberabb committed
337
338

        pbar.close()
339
340
        # reorder all group of results back to original unsorted form
        return re_ords.get_original(res)
baberabb's avatar
baberabb committed
341
342

    def _loglikelihood_tokens(
baberabb's avatar
baberabb committed
343
344
345
        self,
        requests: List[Tuple[Tuple[str, str], List[int], List[int]]],
        disable_tqdm: bool = False,
baberabb's avatar
baberabb committed
346
347
348
349
350
351
352
    ) -> List[Tuple[float, bool]]:
        res = []

        def _collate(x):
            toks = x[1] + x[2]
            return -len(toks), tuple(toks)

353
354
355
356
        # Reorder requests by length and batch
        re_ord = Collator(requests, sort_fn=_collate)
        chunks = re_ord.get_batched(
            n=int(self.batch_size) if self.batch_size != "auto" else 0, batch_fn=None
baberabb's avatar
baberabb committed
357
        )
358

baberabb's avatar
baberabb committed
359
        pbar = tqdm(total=len(requests), disable=disable_tqdm)
baberabb's avatar
baberabb committed
360
        for chunk in chunks:
361
            inputs = []
baberabb's avatar
baberabb committed
362
363
364
365
366
367
368
            ctxlens = []
            for cache_key, context_enc, continuation_enc in chunk:
                inp = (context_enc + continuation_enc)[-(self.max_length) :]
                ctxlen = len(context_enc) - max(
                    0, len(context_enc) + len(continuation_enc) - (self.max_length)
                )

369
                inputs.append(inp)
baberabb's avatar
baberabb committed
370
371
                ctxlens.append(ctxlen)

372
            outputs = self._model_generate(requests=inputs, generate=False)
baberabb's avatar
baberabb committed
373

374
375
            for output, ctxlen, (cache_key, _, _), inp in zip(
                outputs, ctxlens, chunk, inputs
baberabb's avatar
baberabb committed
376
377
            ):
                answer = self._parse_logprobs(
378
379
380
                    tokens=inp,
                    outputs=output,
                    ctxlen=ctxlen,
baberabb's avatar
baberabb committed
381
382
383
384
385
386
387
                )

                res.append(answer)

                # partial caching
                if cache_key is not None:
                    self.cache_hook.add_partial("loglikelihood", cache_key, answer)
388
                pbar.update(1)
baberabb's avatar
baberabb committed
389
390
391
392
        pbar.close()
        return re_ord.get_original(res)

    @staticmethod
baberabb's avatar
baberabb committed
393
    def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]:
baberabb's avatar
baberabb committed
394
395
396
        """Process logprobs and tokens.

        :param tokens: list
397
            Input tokens (potentially left-truncated)
baberabb's avatar
bugfix  
baberabb committed
398
        :param outputs: RequestOutput
399
            Contains prompt_logprobs
baberabb's avatar
baberabb committed
400
401
402
403
404
405
406
407
408
        :param ctxlen: int
            Length of context (so we can slice them away and only keep the predictions)
        :return:
            continuation_logprobs: float
                Log probabilities of continuation tokens
            is_greedy: bool
                Whether argmax matches given continuation exactly
        """

409
        # The first entry of prompt_logprobs is None because the model has no previous tokens to condition on.
baberabb's avatar
bugfix  
baberabb committed
410
411
        continuation_logprobs_dicts = outputs.prompt_logprobs

baberabb's avatar
baberabb committed
412
        # Calculate continuation_logprobs
413
        # assume ctxlen always >= 1
baberabb's avatar
baberabb committed
414
        continuation_logprobs = sum(
baberabb's avatar
baberabb committed
415
            logprob_dict.get(token)
baberabb's avatar
baberabb committed
416
            for token, logprob_dict in zip(
baberabb's avatar
bugfix  
baberabb committed
417
                tokens[ctxlen:], continuation_logprobs_dicts[ctxlen:]
baberabb's avatar
baberabb committed
418
419
420
421
422
            )
        )

        # Determine if is_greedy
        is_greedy = True
baberabb's avatar
baberabb committed
423
424
425
        for token, logprob_dict in zip(
            tokens[ctxlen:], continuation_logprobs_dicts[ctxlen:]
        ):
baberabb's avatar
bugfix  
baberabb committed
426
427
428
429
430
431
            # Get the token with the maximum log probability from the logprob_dict
            if logprob_dict:  # Ensure the logprob_dict is not None
                top_token = max(logprob_dict, key=logprob_dict.get)
                if top_token != token:
                    is_greedy = False
                    break
baberabb's avatar
baberabb committed
432
433

        return continuation_logprobs, is_greedy
434
435
436
437

    @staticmethod
    def modify_gen_kwargs(kwargs: dict) -> dict:
        # sampling_params
438
439
        do_sample = kwargs.pop("do_sample", None)
        if do_sample is False or "temperature" not in kwargs:
440
441
442
443
444
445
446
            kwargs["temperature"] = 0.0
        # hf defaults
        kwargs["skip_special_tokens"] = kwargs.get("skip_special_tokens", False)
        kwargs["spaces_between_special_tokens"] = kwargs.get(
            "spaces_between_special_tokens", False
        )
        return kwargs