From d5246e04c679ce16d10ca92a5f0fd2c080432711 Mon Sep 17 00:00:00 2001 From: sdivyanshu90 Date: Sun, 18 Oct 2026 04:39:19 +0000 Subject: [PATCH] fix(trtllm): respect the generation input limit --- lm_eval/models/trtllm_causallms.py | 60 +++++++++-------- tests/models/test_trtllm_input_limit.py | 88 +++++++++++++++++++++++++ 2 files changed, 122 insertions(+), 26 deletions(-) create mode 100644 tests/models/test_trtllm_input_limit.py diff --git a/lm_eval/models/trtllm_causallms.py b/lm_eval/models/trtllm_causallms.py index dedf62c83a..8321340294 100644 --- a/lm_eval/models/trtllm_causallms.py +++ b/lm_eval/models/trtllm_causallms.py @@ -2,10 +2,11 @@ import logging from importlib.util import find_spec -from typing import List, Literal, Optional, Tuple, Union, Any +from typing import Any, Literal from tqdm import tqdm +from lm_eval.api.instance import Instance from lm_eval.api.model import TemplateLM from lm_eval.api.registry import register_model from lm_eval.models.utils import ( @@ -22,7 +23,7 @@ get_rolling_token_windows, make_disjoint_window, ) -from lm_eval.api.instance import Instance + try: from tensorrt_llm import LLM, SamplingParams @@ -46,9 +47,9 @@ def __init__( model: str, dtype: Literal["float16", "bfloat16", "float32", "auto"] = "auto", trust_remote_code: bool = False, - tokenizer: Optional[str] = None, + tokenizer: str | None = None, tokenizer_mode: Literal["auto", "slow"] = "auto", - add_bos_token: Optional[bool] = None, + add_bos_token: bool | None = None, prefix_token_id: int | None = None, max_batch_size: int = 32, max_input_len: int = 2048, @@ -56,11 +57,11 @@ def __init__( max_beam_width: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, - batch_size: Union[str, int] = "auto", + batch_size: str | int = "auto", max_gen_toks: int = 256, seed: int = 1234, truncation_side: Literal["left", "right", "middle"] = "left", - chat_template_args: Optional[dict] = None, + chat_template_args: dict | None = None, # TRT-LLM: enable thinking tags in apply_chat_template. enable_thinking: bool = True, # End marker for thinking tags - splits to get response after this token (if provided). @@ -106,7 +107,7 @@ def __init__( from transformers import AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained( - tokenizer if tokenizer else model, + tokenizer or model, trust_remote_code=trust_remote_code, use_fast=(tokenizer_mode == "auto"), ) @@ -122,10 +123,12 @@ def __init__( test_tok = self.tok_encode("test", add_special_tokens=True) test_tok_no_special = self.tok_encode("test", add_special_tokens=False) self.add_bos_token = test_tok != test_tok_no_special - except Exception: + except Exception: # noqa: BLE001 - retain fallback for tokenizer-specific failures self.add_bos_token = False - eval_logger.info(f"BOS token will{' not' if not self.add_bos_token else ''} be added") + eval_logger.info( + f"BOS token will{' not' if not self.add_bos_token else ''} be added" + ) if prefix_token_id is not None: eval_logger.info( @@ -238,9 +241,9 @@ def apply_chat_template( def _loglikelihood_tokens( self, - requests: List[Tuple[Tuple[str, str], List[int], List[int]]], + requests: list[tuple[tuple[str, str], list[int], list[int]]], disable_tqdm: bool = False, - ) -> List[Tuple[float, bool]]: + ) -> list[tuple[float, bool]]: res = [] def _collate(x): @@ -260,7 +263,7 @@ def _collate(x): for chunk in chunks: inputs = [] ctxlens = [] - for cache_key, context_enc, continuation_enc in chunk: + 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) @@ -275,10 +278,10 @@ def _collate(x): max_tokens=1, temperature=0, prompt_logprobs=1, - ) + ), ) for output, ctxlen, (cache_key, _, _), inp in zip( - outputs, ctxlens, chunk, inputs + outputs, ctxlens, chunk, inputs, strict=False ): answer = self._parse_logprobs( tokens=inp, @@ -297,7 +300,7 @@ def _collate(x): return re_ord.get_original(res) @staticmethod - def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]: + def _parse_logprobs(tokens: list, outputs, ctxlen: int) -> tuple[float, bool]: """Process logprobs and tokens. :param tokens: list @@ -331,7 +334,9 @@ def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]: continuation_logprobs_dicts = shifted_prompt_logprobs[ctxlen:] continuation_logprobs = 0.0 - for token, logprob_dict in zip(tokens[ctxlen:], continuation_logprobs_dicts): + for token, logprob_dict in zip( + tokens[ctxlen:], continuation_logprobs_dicts, strict=False + ): if token in logprob_dict: continuation_logprobs += logprob_dict[token].logprob else: @@ -344,7 +349,9 @@ def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]: # Determine is_greedy: the actual token must have rank=1 at every position. # The Logprob object has a .rank attribute (1 = highest probability). is_greedy = True - for token, logprob_dict in zip(tokens[ctxlen:], continuation_logprobs_dicts): + for token, logprob_dict in zip( + tokens[ctxlen:], continuation_logprobs_dicts, strict=False + ): if token in logprob_dict and logprob_dict[token].rank is not None: if logprob_dict[token].rank != 1: is_greedy = False @@ -412,12 +419,13 @@ def _collate_gen(_requests): gen_kwargs, eos=eos, default_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 + # Respect both the engine's input cap and its total sequence cap. + # Unused output capacity cannot be borrowed by the prompt. toks, max_gen_toks = maybe_truncate( toks, max_gen_toks=max_gen_toks, - max_model_len=self.max_input_len + self.max_output_len, + max_model_len=self.max_input_len + + min(self.max_output_len, max_gen_toks), side=self.truncation_side, verbose=True, ) @@ -460,8 +468,8 @@ def loglikelihood_rolling( self, requests, disable_tqdm: bool = False, - override_bs: Optional[int] = None, - ) -> List[float]: + override_bs: int | None = None, + ) -> list[float]: """Compute perplexity using rolling windows. This method efficiently batches all rolling windows from all requests together @@ -494,7 +502,7 @@ def loglikelihood_rolling( # Create rolling windows using the utility function # This handles the windowing logic properly with context - rolling_token_windows: List[Tuple[List[int], List[int]]] = list( + rolling_token_windows: list[tuple[list[int], list[int]]] = list( map( make_disjoint_window, get_rolling_token_windows( @@ -523,7 +531,7 @@ def loglikelihood_rolling( 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) + batch_indices, batch_windows = zip(*batch, strict=False) # Process the batch using _loglikelihood_tokens batch_nlls = self._loglikelihood_tokens( @@ -532,7 +540,7 @@ def loglikelihood_rolling( ) # Store results with their request indices - all_nlls.extend(zip(batch_indices, batch_nlls)) + all_nlls.extend(zip(batch_indices, batch_nlls, strict=False)) # Reconstruct per-request loglikelihoods loglikelihoods = [] @@ -587,7 +595,7 @@ def vocab_size(self) -> int: def __del__(self): """Cleanup TRT-LLM resources.""" - if hasattr(self, 'llm'): + if hasattr(self, "llm"): del self.llm @staticmethod diff --git a/tests/models/test_trtllm_input_limit.py b/tests/models/test_trtllm_input_limit.py new file mode 100644 index 0000000000..5c3ac1e9c7 --- /dev/null +++ b/tests/models/test_trtllm_input_limit.py @@ -0,0 +1,88 @@ +"""CPU coverage of TensorRT-LLM's separate input and sequence limits.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import transformers + +from lm_eval.models import trtllm_causallms as trtllm + + +class CharacterTokenizer: + """Use one token per character without downloads or special tokens.""" + + pad_token_id = 0 + eos_token_id = 0 + bos_token_id = 1 + + def decode(self, token): + return "" if token == 1 else "" + + def __call__(self, strings, **kwargs): + return SimpleNamespace(input_ids=[[ord(char) for char in s] for s in strings]) + + +@pytest.fixture +def model(monkeypatch): + """Run the real constructor with a local tokenizer and a recording engine.""" + engine = SimpleNamespace( + generate=Mock( + side_effect=lambda inputs, **kwargs: [ + SimpleNamespace(outputs=[SimpleNamespace(text="answer")]) + for _ in inputs + ] + ) + ) + factory = Mock(return_value=engine) + monkeypatch.setattr(trtllm, "find_spec", lambda name: True) + monkeypatch.setattr(trtllm, "LLM", factory, raising=False) + monkeypatch.setattr(trtllm, "SamplingParams", SimpleNamespace, raising=False) + monkeypatch.setattr(trtllm, "KvCacheConfig", SimpleNamespace, raising=False) + monkeypatch.setattr( + transformers.AutoTokenizer, + "from_pretrained", + Mock(return_value=CharacterTokenizer()), + ) + model = trtllm.TRTLLM( + model="unused", + max_input_len=8, + max_output_len=4, + max_gen_toks=2, + add_bos_token=False, + enable_thinking=False, + ) + assert factory.call_args.kwargs["max_input_len"] == 8 + assert factory.call_args.kwargs["max_seq_len"] == 12 + return model + + +@pytest.mark.parametrize( + ("side", "expected"), + [("left", "mnopqrst"), ("right", "abcdefgh"), ("middle", "abcdqrst")], +) +def test_long_prompt_respects_input_limit(model, side, expected): + """Unused output capacity must not increase the permitted input length.""" + model.truncation_side = side + request = SimpleNamespace(args=("abcdefghijklmnopqrst", {"until": ["\n"]})) + assert model.generate_until([request], disable_tqdm=True) == ["answer"] + call = model.llm.generate.call_args + assert call.kwargs["inputs"] == [[ord(char) for char in expected]] + assert call.kwargs["sampling_params"][0].max_tokens == 2 + + +@pytest.mark.parametrize("prompt_length", [3, 8, 9, 20]) +@pytest.mark.parametrize("generation_length", [1, 4, 6]) +def test_input_and_sequence_limits(model, prompt_length, generation_length): + """Keep short prompts and generation budgets while satisfying both caps.""" + request = SimpleNamespace( + args=("a" * prompt_length, {"max_gen_toks": generation_length, "until": ["\n"]}) + ) + assert model.generate_until([request], disable_tqdm=True) == ["answer"] + call = model.llm.generate.call_args + tokens = call.kwargs["inputs"][0] + params = call.kwargs["sampling_params"][0] + assert len(tokens) <= 8 + assert len(tokens) + params.max_tokens <= 12 + assert len(tokens) == min(prompt_length, 8, 12 - generation_length) + assert params.max_tokens == generation_length