vllm_causallms.py 31.2 KB
Newer Older
Baber's avatar
types  
Baber committed
1
2
from __future__ import annotations

3
import copy
4
import gc
Lintang Sutawika's avatar
Lintang Sutawika committed
5
import logging
6
import os
Baber Abbasi's avatar
Baber Abbasi committed
7
from importlib.metadata import version
8
from importlib.util import find_spec
9
10
11
from multiprocessing import Process, Queue
from queue import Empty
from time import sleep
Baber's avatar
types  
Baber committed
12
from typing import TYPE_CHECKING, Literal
13

14
import jinja2
15
from more_itertools import distribute
Baber Abbasi's avatar
Baber Abbasi committed
16
from packaging.version import parse as parse_version
17
18
from tqdm import tqdm

baberabb's avatar
baberabb committed
19
from lm_eval.api.instance import Instance
20
from lm_eval.api.model import TemplateLM
baberabb's avatar
baberabb committed
21
from lm_eval.api.registry import register_model
22
23
24
25
from lm_eval.models.utils import (
    Collator,
    configure_pad_token,
    handle_stop_sequences,
26
    postprocess_generated_text,
27
28
    undistribute,
)
29
30
31
32
from lm_eval.utils import (
    get_rolling_token_windows,
    make_disjoint_window,
)
33

Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
34

35
try:
36
    import ray
37
    from vllm import LLM, SamplingParams, TokensPrompt
38
    from vllm.lora.request import LoRARequest
baberabb's avatar
baberabb committed
39
    from vllm.transformers_utils.tokenizer import get_tokenizer
40
    from vllm.utils import get_open_port
41
42
43

    if parse_version(version("vllm")) >= parse_version("0.8.3"):
        from vllm.entrypoints.chat_utils import resolve_hf_chat_template
44
except ModuleNotFoundError:
Baber's avatar
Baber committed
45
    print("njklsfnljnlsjnjlksnljnfvljnflsdnlksfnlkvnlksfvnlsfd")
Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
46

47
48
if TYPE_CHECKING:
    pass
bcicc's avatar
bcicc committed
49

Lintang Sutawika's avatar
Lintang Sutawika committed
50
eval_logger = logging.getLogger(__name__)
baberabb's avatar
baberabb committed
51

baberabb's avatar
baberabb committed
52

53
54
def _vllm_mp_worker(
    model_args: dict,
fxmarty-amd's avatar
fxmarty-amd committed
55
    sampling_params: list["SamplingParams"],
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
    requests: list[list[int]],
    lora_request: "LoRARequest",
    result_queue: "Queue",
    dp_size: int,
    local_dp_rank: int,
    dp_master_port: int,
    dp_master_ip: str = "127.0.0.1",
) -> None:
    """
    Worker process for vLLM multiprocessing.
    Initializes a vLLM engine, processes requests, and puts results or errors
    onto the result_queue.
    """

    if not requests:
        result_queue.put((local_dp_rank, []))
        return None

    os.environ["VLLM_DP_RANK"] = os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
    os.environ["VLLM_DP_SIZE"] = str(dp_size)
    os.environ["VLLM_DP_MASTER_IP"] = str(dp_master_ip)
    os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)

    llm = None
    try:
        llm = LLM(**model_args)
        res = llm.generate(
83
            [TokensPrompt(prompt_token_ids=request) for request in requests],
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
            sampling_params=sampling_params,
            lora_request=lora_request,
        )
        # Give engines time to pause their processing loops before exiting."
        sleep(1)
        result_queue.put((local_dp_rank, res))

    except Exception as e:
        error_message = f"Worker {local_dp_rank} failed during generation: {type(e).__name__}: {str(e)}"
        eval_logger.error(error_message, exc_info=True)
        result_queue.put((local_dp_rank, {"error": error_message}))

    finally:
        if llm is not None:
            try:
                del llm
                gc.collect()
            except Exception as e_cleanup:
                eval_logger.warning(
                    f"Worker {local_dp_rank} encountered an error during LLM cleanup: {type(e_cleanup).__name__}: {str(e_cleanup)}",
                    exc_info=True,
                )

    return None


baberabb's avatar
baberabb committed
110
@register_model("vllm")
111
class VLLM(TemplateLM):
baberabb's avatar
baberabb committed
112
113
114
115
    _DEFAULT_MAX_LENGTH = 2048

    def __init__(
        self,
116
        pretrained: str,
baberabb's avatar
baberabb committed
117
        dtype: Literal["float16", "bfloat16", "float32", "auto"] = "auto",
Baber's avatar
types  
Baber committed
118
119
120
        revision: str | None = None,
        trust_remote_code: bool | None = False,
        tokenizer: str | None = None,
baberabb's avatar
baberabb committed
121
        tokenizer_mode: Literal["auto", "slow"] = "auto",
Baber's avatar
types  
Baber committed
122
123
124
        tokenizer_revision: str | None = None,
        add_bos_token: bool | None = False,
        prefix_token_id: int | None = None,
baberabb's avatar
baberabb committed
125
        tensor_parallel_size: int = 1,
Baber's avatar
types  
Baber committed
126
        quantization: str | None = None,
baberabb's avatar
baberabb committed
127
128
        max_gen_toks: int = 256,
        swap_space: int = 4,
Baber's avatar
types  
Baber committed
129
130
131
132
        batch_size: str | int = 1,
        max_batch_size: int | None = None,
        max_length: int | None = None,
        max_model_len: int | None = None,
baberabb's avatar
baberabb committed
133
        seed: int = 1234,
134
        gpu_memory_utilization: float = 0.9,
135
        data_parallel_size: int = 1,
Baber's avatar
types  
Baber committed
136
        lora_local_path: str | None = None,
137
138
        # VLLM: enable thinking tags in the prompt.
        enable_thinking: bool = True,
139
        chat_template_args: Optional[dict] = None,
140
        # End marker for thinking tags - splits to get response after this token (if provided).
Baber's avatar
types  
Baber committed
141
        think_end_token: str | None = None,
MaYongQing's avatar
MaYongQing committed
142
        max_lora_rank: int = 16,
Baber Abbasi's avatar
Baber Abbasi committed
143
        **kwargs,
baberabb's avatar
baberabb committed
144
145
    ):
        super().__init__()
146

147
        if not find_spec("vllm"):
148
            raise ModuleNotFoundError(
149
150
                "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
151
152
            )

Baber Abbasi's avatar
Baber Abbasi committed
153
154
155
        assert max_length is None or max_model_len is None, (
            "Either max_length or max_model_len may be provided, but not both"
        )
Baber Abbasi's avatar
Baber Abbasi committed
156
        kwargs.pop("device", None)
157
        self.think_end_token = think_end_token
158
        self.V1 = os.environ.get("VLLM_USE_V1", "1") != "0"
159
        self._max_length = max_model_len if max_model_len is not None else max_length
baberabb's avatar
baberabb committed
160
        self.tensor_parallel_size = int(tensor_parallel_size)
161
        self.data_parallel_size = int(data_parallel_size)
baberabb's avatar
baberabb committed
162
163
164
165
166
        self.model_args = {
            "model": pretrained,
            "gpu_memory_utilization": float(gpu_memory_utilization),
            "revision": revision,
            "dtype": dtype,
baberabb's avatar
baberabb committed
167
            "tokenizer": tokenizer,
baberabb's avatar
baberabb committed
168
            "tokenizer_mode": tokenizer_mode,
baberabb's avatar
baberabb committed
169
            "tokenizer_revision": tokenizer_revision,
baberabb's avatar
baberabb committed
170
171
            "trust_remote_code": trust_remote_code,
            "tensor_parallel_size": int(tensor_parallel_size),
172
            "max_model_len": int(self._max_length) if self._max_length else None,
173
            "max_num_seqs": kwargs.get("max_num_seqs", max_batch_size),
baberabb's avatar
baberabb committed
174
175
176
            "swap_space": int(swap_space),
            "quantization": quantization,
            "seed": int(seed),
Baber's avatar
types  
Baber committed
177
            "enable_lora": bool(lora_local_path),
MaYongQing's avatar
MaYongQing committed
178
            "max_lora_rank": int(max_lora_rank),
baberabb's avatar
baberabb committed
179
        }
Baber Abbasi's avatar
Baber Abbasi committed
180
        self.model_args.update(kwargs)
181
182
183
        self.batch_size = (
            "auto"
            if isinstance(batch_size, str) and "auto" in batch_size
184
            else int(batch_size)
185
        )
186
        if self.data_parallel_size <= 1:
baberabb's avatar
baberabb committed
187
            self.model = LLM(**self.model_args)
baberabb's avatar
baberabb committed
188
        else:
Baber Abbasi's avatar
Baber Abbasi committed
189
190
191
            eval_logger.warning(
                "You might experience occasional issues with model weight downloading when data_parallel is in use. To ensure stable performance, run with data_parallel_size=1 until the weights are downloaded and cached."
            )
192
193
194
195
196
            self.model_args["distributed_executor_backend"] = (
                "ray"
                if not self.V1
                else self.model_args.get("distributed_executor_backend", None)
            )
197
198
199
            self.batch_size = "auto"
            eval_logger.info("Manual batching is not compatible with data parallelism.")

200
201
202
203
204
205
        if "gemma" in pretrained.lower():
            add_bos_token = True
            eval_logger.info(
                "Found 'gemma' in model name, a BOS token will be used as Gemma series models underperform without it."
            )

206
        from transformers import AutoConfig
207

208
209
210
        self._config = AutoConfig.from_pretrained(
            pretrained, trust_remote_code=trust_remote_code, revision=revision
        )
baberabb's avatar
nits  
baberabb committed
211
212
213
214
        self.tokenizer = get_tokenizer(
            tokenizer if tokenizer else pretrained,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=trust_remote_code,
215
            revision=tokenizer_revision,
216
            add_bos_token=add_bos_token,
baberabb's avatar
nits  
baberabb committed
217
        )
218
        self.tokenizer = configure_pad_token(self.tokenizer, model_config=self._config)
219
        self.chat_template_args = chat_template_args or {}
220
        self.enable_thinking = self.chat_template_args.pop(
221
222
            "enable_thinking", enable_thinking
        )
223
        self.add_bos_token = add_bos_token
224

225
        if parse_version(version("vllm")) >= parse_version("0.8.3"):
226
227
228
229
230
231
232
            kwargs_resolve_hf_chat_template = {
                "tokenizer": self.tokenizer,
                "chat_template": None,
                "tools": None,
            }

            if parse_version(version("vllm")) >= parse_version("0.9.0"):
233
234
235
236
237
238
239
240
241
242
243
                if self.data_parallel_size <= 1:
                    kwargs_resolve_hf_chat_template["model_config"] = (
                        self.model.llm_engine.model_config
                    )
                else:
                    from vllm.engine.arg_utils import EngineArgs

                    engine_args = EngineArgs(**self.model_args)
                    model_config = engine_args.create_model_config()

                    kwargs_resolve_hf_chat_template["model_config"] = model_config
244
245
246
            else:
                kwargs_resolve_hf_chat_template["trust_remote_code"] = trust_remote_code

247
            self.hf_chat_template = resolve_hf_chat_template(
248
                **kwargs_resolve_hf_chat_template
249
250
251
            )
        else:
            self.hf_chat_template = None
252

253
254
255
256
257
        self.custom_prefix_token_id = prefix_token_id
        if prefix_token_id is not None:
            eval_logger.info(
                f"Loglikelihood prefix token id used in evaluation: {self.prefix_token_id}"
            )
258

baberabb's avatar
baberabb committed
259
260
        self._max_gen_toks = max_gen_toks

bcicc's avatar
bcicc committed
261
        if lora_local_path is not None:
Baber Abbasi's avatar
Baber Abbasi committed
262
263
264
            assert parse_version(version("vllm")) > parse_version("0.3.0"), (
                "lora adapters only compatible with vllm > v0.3.0."
            )
bcicc's avatar
bcicc committed
265
266
267
268
            self.lora_request = LoRARequest("finetuned", 1, lora_local_path)
        else:
            self.lora_request = None

baberabb's avatar
baberabb committed
269
270
271
272
273
    @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

274
275
276
277
278
279
280
281
282
    @property
    def prefix_token_id(self):
        # it is used as prefix for loglikelihood
        if self.custom_prefix_token_id is not None:
            return self.custom_prefix_token_id
        if self.tokenizer.bos_token_id is not None:
            return self.tokenizer.bos_token_id
        return self.tokenizer.eos_token_id

baberabb's avatar
baberabb committed
283
284
285
286
    @property
    def max_length(self):
        if self._max_length:  # if max length manually set, return it
            return self._max_length
287
288
289
290
291
292
293
294
295
296
297
298
        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
299
300
301
302
303

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

Baber Abbasi's avatar
Baber Abbasi committed
304
    def apply_chat_template(
Baber's avatar
types  
Baber committed
305
        self, chat_history: list[dict[str, str]], add_generation_prompt: bool = True
Baber Abbasi's avatar
Baber Abbasi committed
306
    ) -> str:
307
308
309
        """
        Method to apply a chat template to a list of chat history between user and model.
        """
310
311
312
313
314
315
316
317
        try:
            chat_templated = self.tokenizer.apply_chat_template(
                chat_history,
                tokenize=False,
                add_generation_prompt=add_generation_prompt,
                continue_final_message=not add_generation_prompt,
                chat_template=self.hf_chat_template,
                enable_thinking=self.enable_thinking,
318
                **self.chat_template_args,
319
320
321
322
323
324
325
326
327
328
329
330
            )
        except jinja2.exceptions.TemplateError:
            eval_logger.warning(
                "Failed to apply chat template. removing the system role in chat history."
            )
            chat_templated = self.tokenizer.apply_chat_template(
                [msg for msg in chat_history if msg["role"] != "system"],
                tokenize=False,
                add_generation_prompt=add_generation_prompt,
                continue_final_message=not add_generation_prompt,
                chat_template=self.hf_chat_template,
                enable_thinking=self.enable_thinking,
331
                **self.chat_template_args,
332
            )
333

Baber Abbasi's avatar
Baber Abbasi committed
334
335
        return chat_templated

336
337
338
339
    @property
    def tokenizer_name(self) -> str:
        return self.tokenizer.name_or_path.replace("/", "__")

baberabb's avatar
baberabb committed
340
341
    def tok_encode(
        self,
Baber's avatar
types  
Baber committed
342
        string: str | list[str],
343
344
345
        left_truncate_len: int = None,
        add_special_tokens: bool = False,
        truncation: bool = False,
Baber's avatar
types  
Baber committed
346
    ) -> list[int] | list[list[int]]:
347
348
        if not add_special_tokens:
            add_special_tokens = False or self.add_bos_token
Baber's avatar
types  
Baber committed
349
        encoding: list[list[int]] | list[int] = self.tokenizer(
350
351
352
353
354
            string,
            add_special_tokens=add_special_tokens,
            truncation=truncation,
            return_attention_mask=False,
        ).input_ids
baberabb's avatar
baberabb committed
355
356
357

        # left-truncate the encoded context to be at most `left_truncate_len` tokens long
        if left_truncate_len:
358
359
360
361
            if not isinstance(string, str):
                encoding = [enc[-left_truncate_len:] for enc in encoding]
            else:
                encoding = encoding[-left_truncate_len:]
baberabb's avatar
baberabb committed
362
363
364
365
366

        return encoding

    def _model_generate(
        self,
Baber's avatar
types  
Baber committed
367
        requests: list[list[int]] = None,
baberabb's avatar
baberabb committed
368
        generate: bool = False,
fxmarty-amd's avatar
fxmarty-amd committed
369
        sampling_params: Union[List["SamplingParams"], "SamplingParams", None] = None,
baberabb's avatar
baberabb committed
370
    ):
371
        if not generate or sampling_params is None:
baberabb's avatar
baberabb committed
372
            sampling_params = SamplingParams(
373
                temperature=0, prompt_logprobs=1, max_tokens=1, detokenize=False
baberabb's avatar
baberabb committed
374
            )
375
376
        if not isinstance(sampling_params, List):
            sampling_params = [sampling_params] * len(requests)
377
        if self.data_parallel_size > 1 and not self.V1:
Baber Abbasi's avatar
Baber Abbasi committed
378
            # vLLM hangs if resources are set in ray.remote
Baber Abbasi's avatar
Baber Abbasi committed
379
380
            # also seems to only work with decorator and not with ray.remote() fn
            # see https://github.com/vllm-project/vllm/issues/973
Baber Abbasi's avatar
Baber Abbasi committed
381
            @ray.remote
Baber Abbasi's avatar
Baber Abbasi committed
382
            def run_inference_one_model(
383
                model_args: dict,
Baber's avatar
types  
Baber committed
384
385
                sampling_params: list["SamplingParams"],
                requests: list[list[int]],
fxmarty-amd's avatar
fxmarty-amd committed
386
                lora_request: "LoRARequest",
Baber Abbasi's avatar
Baber Abbasi committed
387
388
389
            ):
                llm = LLM(**model_args)
                return llm.generate(
390
                    [TokensPrompt(prompt_token_ids=request) for request in requests],
391
392
                    sampling_params=sampling_params,
                    lora_request=lora_request,
Baber Abbasi's avatar
Baber Abbasi committed
393
394
                )

395
396
397
            # dispatch requests to all self.data_parallel_size workers, in interleaved fashion
            # interleaved important to balance context lengths across workers
            requests = [list(x) for x in distribute(self.data_parallel_size, requests)]
398
399
400
            sampling_params = [
                list(sp) for sp in distribute(self.data_parallel_size, sampling_params)
            ]
401
            inputs = (
402
403
                (self.model_args, sp, req, self.lora_request)
                for req, sp in zip(requests, sampling_params)
404
            )
Baber Abbasi's avatar
Baber Abbasi committed
405
406
            object_refs = [run_inference_one_model.remote(*x) for x in inputs]
            results = ray.get(object_refs)
407
408
            # Invoke ray.shutdown() to prevent hang-ups if subsequent calls required.
            ray.shutdown()
baberabb's avatar
baberabb committed
409
            # flatten results
410
            return undistribute(results)
411
412
413
414
415
416
417
        elif self.data_parallel_size > 1:
            # based on https://github.com/vllm-project/vllm/blob/a04720bc36401d831cb048c3917b9e58173d9c1d/examples/offline_inference/data_parallel.py
            dp_size = self.data_parallel_size
            dp_master_ip = os.environ.get("VLLM_DP_MASTER_IP", "127.0.0.1")
            dp_master_port = os.environ.get("VLLM_DP_MASTER_PORT") or get_open_port()

            requests = (list(x) for x in distribute(self.data_parallel_size, requests))
418
419
420
            sampling_params = (
                list(sp) for sp in distribute(self.data_parallel_size, sampling_params)
            )
421
422
423
            procs, resq = [], Queue()
            # We use Process as it is non-daemonic
            try:
424
                for rank, (sp, req) in enumerate(zip(requests, sampling_params)):
425
426
427
428
                    proc = Process(
                        target=_vllm_mp_worker,
                        args=(
                            self.model_args.copy(),
429
                            sp,
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
                            req,
                            self.lora_request,
                            resq,
                            dp_size,
                            rank,
                            dp_master_port,
                            dp_master_ip,
                        ),
                    )
                    proc.start()
                    procs.append(proc)

                # Collect results
                rank_res = {}
                while len(rank_res) < len(procs):
                    try:
                        rank, result = resq.get(timeout=30)
                        if isinstance(result, dict) and "error" in result:
                            raise RuntimeError(result["error"])
                        rank_res[rank] = result
                    except Empty:
                        dead_procs = [
                            idx
                            for idx, p in enumerate(procs)
                            if not p.is_alive() and idx not in rank_res
                        ]
                        if dead_procs:
                            raise RuntimeError(
                                f"Worker processes {dead_procs} died unexpectedly"
Baber's avatar
types  
Baber committed
459
                            ) from None
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
                        continue

                results = [rank_res[i] for i in range(len(procs))]
                return undistribute(results)

            # cleanup
            finally:
                try:
                    resq.close()
                    resq.join_thread()
                except Exception:
                    eval_logger.debug(
                        "Failed to close vllm DP results queue", exc_info=True
                    )
                for proc in procs:
                    proc.join(timeout=10)
                    if proc.is_alive():
                        proc.terminate()
                        proc.join(timeout=5)
                        if proc.is_alive():
                            proc.kill()
baberabb's avatar
baberabb committed
481

482
483
        else:
            outputs = self.model.generate(
484
                [TokensPrompt(prompt_token_ids=request) for request in requests],
485
                sampling_params=sampling_params,
Baber's avatar
types  
Baber committed
486
                use_tqdm=self.batch_size == "auto",
487
488
489
                lora_request=self.lora_request,
            )
            return outputs
baberabb's avatar
baberabb committed
490

491
    def loglikelihood_rolling(
Baber's avatar
types  
Baber committed
492
493
        self, requests: list[Instance], disable_tqdm: bool = False
    ) -> list[float]:
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        adaptive_batch_size = None
        if self.batch_size == "auto":
            adaptive_batch_size = len(requests)

        # First, collect all windows from all requests
        all_windows = []  # List of (request_idx, window) tuples
        request_window_counts = []  # Track number of windows per request

        for req_idx, (string,) in enumerate(
            tqdm(
                [req.args for req in requests],
                disable=(disable_tqdm or (self.rank != 0)),
            )
        ):
Baber's avatar
types  
Baber committed
508
            rolling_token_windows: list[tuple[list[int], list[int]]] = list(
baberabb's avatar
baberabb committed
509
                map(
510
511
                    make_disjoint_window,
                    get_rolling_token_windows(
baberabb's avatar
baberabb committed
512
                        token_list=self.tok_encode(string),
513
514
                        prefix_token=self.prefix_token_id,
                        # max_seq_len - (1 for context)
baberabb's avatar
baberabb committed
515
                        max_seq_len=self.max_length - 1,
baberabb's avatar
baberabb committed
516
517
518
519
520
                        context_len=1,
                    ),
                )
            )

521
522
            # TODO: Right now, we pass single EOT token to the Encoder and the full context to the decoder, in seq2seq case
            windows = [(None,) + x for x in rolling_token_windows]
baberabb's avatar
baberabb committed
523

524
525
526
            # Store windows with their request index
            all_windows.extend((req_idx, window) for window in windows)
            request_window_counts.append(len(windows))
baberabb's avatar
baberabb committed
527

528
529
530
531
532
533
        all_nlls = []
        batch_size = adaptive_batch_size or int(self.batch_size)
        for i in range(0, len(all_windows), batch_size):
            batch = all_windows[i : i + batch_size]
            # Extract just the windows for processing, keeping track of request indices
            batch_indices, batch_windows = zip(*batch)
baberabb's avatar
baberabb committed
534

535
536
537
538
539
540
            batch_nlls = self._loglikelihood_tokens(
                requests=batch_windows,
                disable_tqdm=False,
            )
            # Store results with their request indices
            all_nlls.extend(zip(batch_indices, batch_nlls))
541

542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
        # Reconstruct per-request loglikelihoods
        loglikelihoods = []
        current_idx = 0
        for window_count in request_window_counts:
            # Get all nlls for this request
            request_nlls = all_nlls[current_idx : current_idx + window_count]
            # Sum up the nlls for this request (discarding is_greedy)
            request_total = sum(nll[0] for _, nll in request_nlls)
            loglikelihoods.append(request_total)
            current_idx += window_count

            string = requests[len(loglikelihoods) - 1].args[0]
            self.cache_hook.add_partial(
                "loglikelihood_rolling", (string,), request_total
            )
557

baberabb's avatar
baberabb committed
558
559
        return loglikelihoods

560
    def generate_until(
Baber's avatar
types  
Baber committed
561
562
        self, requests: list[Instance], disable_tqdm: bool = False
    ) -> list[str]:
563
        res = []
baberabb's avatar
baberabb committed
564
565
566

        # batch tokenize contexts
        context, all_gen_kwargs = zip(*(req.args for req in requests))
Baber's avatar
types  
Baber committed
567
        context_encoding: list[list[int]] = self.tok_encode(
568
569
            context, add_special_tokens=self.add_bos_token
        )
baberabb's avatar
baberabb committed
570
571
572
        requests = [
            ((a, b), c) for a, b, c in zip(context, context_encoding, all_gen_kwargs)
        ]
baberabb's avatar
baberabb committed
573
574
575
576
577
578
579
580

        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
581
            return -len(_requests[0][1]), _requests[0][0]
baberabb's avatar
baberabb committed
582

583
584
585
586
587
        re_ords = Collator(
            requests,
            _collate_gen,
            group_by=None,
        )
588
589
590
        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
591

592
593
        pbar = tqdm(
            total=len(requests),
594
            disable=(disable_tqdm or (self.rank != 0)),
595
596
            desc="Running generate_until requests",
        )
baberabb's avatar
baberabb committed
597
        # for each different set of kwargs, we execute all requests, by batch.
598
        eos = self.tokenizer.decode(self.eot_token_id)
599
600
601
        for chunk in chunks:
            context_and_encoding, all_gen_kwargs = zip(*chunk)
            context, context_encoding = zip(*context_and_encoding)
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
            context_encoding_truncated = []
            sampling_params = []
            for x, gen_kwargs in zip(context_encoding, all_gen_kwargs):
                # unpack our keyword arguments.
                if isinstance(gen_kwargs, dict):
                    kwargs = copy.deepcopy(gen_kwargs)  # edge case for repeats > 1
                    # add EOS token to stop sequences
                    until = handle_stop_sequences(kwargs.pop("until", None), eos=eos)
                else:
                    raise ValueError(
                        f"Expected `kwargs` to be of type `dict` but got {type(gen_kwargs)}"
                    )
                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
                if len(x) > max_ctx_len:
623
                    eval_logger.warning(
624
                        f"Context length {len(x)} exceeds max length (context + max gen tokens): {max_ctx_len}. Truncating context."
625
                    )
626
627
628
629
630
631
632
633
                    context_encoding_truncated.append(x[-max_ctx_len:])
                else:
                    context_encoding_truncated.append(x)
                # create sampling params
                kwargs = self.modify_gen_kwargs(kwargs)
                sampling_params.append(
                    SamplingParams(max_tokens=max_gen_toks, stop=until, **kwargs)
                )
634
635
636

            # perform batched generation
            cont = self._model_generate(
637
                requests=context_encoding_truncated,
638
                generate=True,
639
                sampling_params=sampling_params,
640
            )
baberabb's avatar
baberabb committed
641

642
            # cache generations
Baber's avatar
types  
Baber committed
643
            for output, context_ in zip(cont, context):
644
                generated_text: str = output.outputs[0].text
645
                # use secondary stop seqs to cut off should-have-been-stopped content post-hoc
646
647
648
                generated_text = postprocess_generated_text(
                    generated_text, until, self.think_end_token
                )
649
650
                res.append(generated_text)
                self.cache_hook.add_partial(
Baber's avatar
types  
Baber committed
651
                    "generate_until", (context_, gen_kwargs), generated_text
652
653
                )
                pbar.update(1)
baberabb's avatar
baberabb committed
654
655

        pbar.close()
656
657
        # reorder all group of results back to original unsorted form
        return re_ords.get_original(res)
baberabb's avatar
baberabb committed
658
659

    def _loglikelihood_tokens(
baberabb's avatar
baberabb committed
660
        self,
Baber's avatar
types  
Baber committed
661
        requests: list[tuple[tuple[str, str], list[int], list[int]]],
baberabb's avatar
baberabb committed
662
        disable_tqdm: bool = False,
Baber's avatar
types  
Baber committed
663
    ) -> list[tuple[float, bool]]:
baberabb's avatar
baberabb committed
664
665
666
667
668
669
        res = []

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

670
671
672
673
        # 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
674
        )
675

676
677
678
679
680
        pbar = tqdm(
            total=len(requests),
            disable=disable_tqdm,
            desc="Running loglikelihood requests",
        )
baberabb's avatar
baberabb committed
681
        for chunk in chunks:
682
            inputs = []
baberabb's avatar
baberabb committed
683
            ctxlens = []
Baber's avatar
types  
Baber committed
684
            for _cache_key, context_enc, continuation_enc in chunk:
685
686
                if (
                    full_length := len(context_enc + continuation_enc)
687
                ) > self.max_length:
688
689
690
                    eval_logger.warning(
                        f"Context length {full_length} exceeds max length ({self.max_length}). Truncating context."
                    )
baberabb's avatar
baberabb committed
691
692
693
694
695
                inp = (context_enc + continuation_enc)[-(self.max_length) :]
                ctxlen = len(context_enc) - max(
                    0, len(context_enc) + len(continuation_enc) - (self.max_length)
                )

696
                inputs.append(inp)
baberabb's avatar
baberabb committed
697
698
                ctxlens.append(ctxlen)

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

701
702
            for output, ctxlen, (cache_key, _, _), inp in zip(
                outputs, ctxlens, chunk, inputs
baberabb's avatar
baberabb committed
703
704
            ):
                answer = self._parse_logprobs(
705
706
707
                    tokens=inp,
                    outputs=output,
                    ctxlen=ctxlen,
baberabb's avatar
baberabb committed
708
709
710
711
712
                )

                res.append(answer)

                if cache_key is not None:
713
714
715
                    # special case: loglikelihood_rolling produces a number of loglikelihood requests
                    # all with cache key None. instead do add_partial on the per-example level
                    # in the loglikelihood_rolling() function for those.
baberabb's avatar
baberabb committed
716
                    self.cache_hook.add_partial("loglikelihood", cache_key, answer)
717
                pbar.update(1)
baberabb's avatar
baberabb committed
718
719
720
721
        pbar.close()
        return re_ord.get_original(res)

    @staticmethod
Baber's avatar
types  
Baber committed
722
    def _parse_logprobs(tokens: list, outputs, ctxlen: int) -> tuple[float, bool]:
baberabb's avatar
baberabb committed
723
724
725
        """Process logprobs and tokens.

        :param tokens: list
726
            Input tokens (potentially left-truncated)
baberabb's avatar
bugfix  
baberabb committed
727
        :param outputs: RequestOutput
728
            Contains prompt_logprobs
baberabb's avatar
baberabb committed
729
730
731
732
733
734
735
736
737
        :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
        """

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

741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
        def coerce_logprob_to_num(logprob):
            # vLLM changed the return type of logprobs from float
            # to a Logprob object storing the float value + extra data
            # (https://github.com/vllm-project/vllm/pull/3065).
            # If we are dealing with vllm's Logprob object, return
            # the logprob value stored as an attribute. Otherwise,
            # return the object itself (which should be a float
            # for older versions of vLLM).
            return getattr(logprob, "logprob", logprob)

        continuation_logprobs_dicts = [
            {
                token: coerce_logprob_to_num(logprob)
                for token, logprob in logprob_dict.items()
            }
            if logprob_dict is not None
            else None
            for logprob_dict in continuation_logprobs_dicts
        ]

baberabb's avatar
baberabb committed
761
        # Calculate continuation_logprobs
762
        # assume ctxlen always >= 1
baberabb's avatar
baberabb committed
763
        continuation_logprobs = sum(
baberabb's avatar
baberabb committed
764
            logprob_dict.get(token)
baberabb's avatar
baberabb committed
765
            for token, logprob_dict in zip(
baberabb's avatar
bugfix  
baberabb committed
766
                tokens[ctxlen:], continuation_logprobs_dicts[ctxlen:]
baberabb's avatar
baberabb committed
767
768
769
770
771
            )
        )

        # Determine if is_greedy
        is_greedy = True
baberabb's avatar
baberabb committed
772
773
774
        for token, logprob_dict in zip(
            tokens[ctxlen:], continuation_logprobs_dicts[ctxlen:]
        ):
baberabb's avatar
bugfix  
baberabb committed
775
776
777
778
779
780
            # 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
781
782

        return continuation_logprobs, is_greedy
783
784
785
786

    @staticmethod
    def modify_gen_kwargs(kwargs: dict) -> dict:
        # sampling_params
787
        kwargs["temperature"] = kwargs.get("temperature", 0.0)
788
        do_sample = kwargs.pop("do_sample", None)
789
790
791
792
        if do_sample is False and "temperature" not in kwargs:
            eval_logger.debug(
                "Got `do_sample=False` and no temperature value, setting VLLM temperature to 0.0 ..."
            )
793
794
795
796
797
798
799
            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