From b68114bbcef999e8ce117b03ba03aa9d5c595f78 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 4 Aug 2026 20:03:28 +0000 Subject: [PATCH 1/7] Use lm-eval's built-in trtllm backend, deprecate lm_eval_tensorrt_llm.py lm-evaluation-harness 0.4.12 is the first release shipping a TensorRT-LLM backend (lm_eval.models.trtllm_causallms, registered as `trtllm`), so the example no longer needs its own. Pin lm_eval>=0.4.12,<0.5 (0.5.0.dev drops the file) and bump lm_eval_hf.py's version guard to match. lm_eval_tensorrt_llm.py becomes a deprecation shim: it warns, translates the legacy arguments (checkpoint_dir -> model, max_length -> max_input_len / max_output_len, --batch_size -> max_batch_size, tensor parallelism over all visible GPUs) and forwards to `lm_eval_hf.py --model trtllm`, so existing commands keep working. The upstream backend accepts **kwargs but only forwards a fixed set to the TRT-LLM LLM API, so two of its defaults have to be set on every command: tensor_parallel_size defaults to 1, and max_input_len defaults to 2048, which silently left-truncates 5-shot prompts. huggingface_example.sh and the docs now pass both explicitly. It also misreads TensorRT-LLM's prompt_logprobs. TRT-LLM aligns them to the next token (base_worker.py passes prompt_token_ids[1:] + first generated token), so entry i predicts tokens[i + 1]; _parse_logprobs reads prompt_logprobs[i][tokens[i]] and shifts on top of that, raising KeyError on the first request of every loglikelihood task. lm_eval_hf.py patches the alignment when --model trtllm is selected; the patch goes away once that is fixed upstream. Tested on nvidia/Qwen3.5-122B-A10B-NVFP4 (4x B300, TRT-LLM 1.3.0rc23) at tp=1/2/4: hellaswag acc 0.7188 / acc_norm 0.7812 at every tp, identical to the removed implementation; gsm8k and the deprecation shim also verified. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../skills/deployment/references/trtllm.md | 18 +- CHANGELOG.rst | 2 + .../hf_ptq/scripts/huggingface_example.sh | 13 +- examples/llm_eval/README.md | 25 +- examples/llm_eval/lm_eval_hf.py | 45 ++- examples/llm_eval/lm_eval_tensorrt_llm.py | 299 +++++++----------- examples/llm_eval/requirements.txt | 2 +- tests/examples/llm_eval/test_llm_eval.py | 5 +- 8 files changed, 200 insertions(+), 209 deletions(-) diff --git a/.agents/skills/deployment/references/trtllm.md b/.agents/skills/deployment/references/trtllm.md index 865c11e7eee..59027774fb1 100644 --- a/.agents/skills/deployment/references/trtllm.md +++ b/.agents/skills/deployment/references/trtllm.md @@ -69,14 +69,20 @@ If you encounter a legacy checkpoint (no `hf_quant_config.json`, has `rank*.safe ## Evaluation with TRT-LLM -```python -# examples/llm_eval/lm_eval_tensorrt_llm.py -# Runs lm_evaluation_harness benchmarks with TRT-LLM -python examples/llm_eval/lm_eval_tensorrt_llm.py \ - --model_path \ - --tasks gsm8k,mmlu +Runs lm-evaluation-harness benchmarks through lm-eval's built-in `trtllm` backend +(requires `lm_eval>=0.4.12`; ModelOpt's own `lm_eval_tensorrt_llm.py` is deprecated). + +```bash +python examples/llm_eval/lm_eval_hf.py \ + --model trtllm \ + --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512 \ + --tasks gsm8k,mmlu \ + --batch_size ``` +`max_input_len` defaults to 2048 and longer prompts are silently truncated, so set it +explicitly for few-shot tasks. + ## Common Issues | Issue | Fix | diff --git a/CHANGELOG.rst b/CHANGELOG.rst index afaec6bc805..30129f067a1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,8 @@ Changelog **Deprecations** +- Deprecate ``examples/llm_eval/lm_eval_tensorrt_llm.py`` (the ``trt-llm`` model). lm-evaluation-harness 0.4.12 ships its own TensorRT-LLM backend (``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), which also implements ``loglikelihood_rolling`` and pipeline parallelism, so the example uses that instead and pins ``lm_eval>=0.4.12,<0.5``. The old entry point still runs: it warns, translates the legacy arguments (``checkpoint_dir`` → ``model``, ``max_length`` → ``max_input_len``/``max_output_len``, ``--batch_size`` → ``max_batch_size``, tensor parallelism over all visible GPUs) and forwards to ``lm_eval_hf.py --model trtllm``. When calling the new backend directly, set ``tensor_parallel_size`` and ``max_input_len`` explicitly — they default to 1 and 2048, and longer prompts are silently truncated. ``lm_eval_hf.py`` also patches that backend's ``_parse_logprobs``, which is off by one against TensorRT-LLM's next-token-aligned ``prompt_logprobs`` and otherwise raises ``KeyError`` on every loglikelihood task; run TensorRT-LLM evaluations through ``lm_eval_hf.py`` rather than the plain ``lm_eval`` CLI until that is fixed upstream. + **Bug Fixes** 0.46 (2026-08-xx) diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index 97924c90ff3..fadb97d0070 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -297,11 +297,16 @@ if [[ $TASKS =~ "lm_eval" ]]; then pip install -r requirements.txt - echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE" + # lm-eval's `trtllm` backend defaults to 1 GPU; shard over every visible one instead. + LM_EVAL_TP=$(python -c "import torch; print(max(torch.cuda.device_count(), 1))") - python lm_eval_tensorrt_llm.py \ - --model trt-llm \ - --model_args tokenizer=$MODEL_PATH,checkpoint_dir=$SAVE_PATH,max_gen_toks=$BUILD_MAX_OUTPUT_LEN \ + echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE tp $LM_EVAL_TP" + + # max_input_len defaults to 2048, which silently truncates 5-shot prompts; size the + # engine for a 4096-token context plus the requested generation length. + python lm_eval_hf.py \ + --model trtllm \ + --model_args model=$SAVE_PATH,tokenizer=$MODEL_PATH,tensor_parallel_size=$LM_EVAL_TP,max_batch_size=$BUILD_MAX_BATCH_SIZE,max_gen_toks=$BUILD_MAX_OUTPUT_LEN,max_input_len=4096,max_output_len=$BUILD_MAX_OUTPUT_LEN \ --tasks $LM_EVAL_TASKS \ --batch_size $BUILD_MAX_BATCH_SIZE $lm_eval_flags | tee $LM_EVAL_RESULT diff --git a/examples/llm_eval/README.md b/examples/llm_eval/README.md index f680235b651..af0dfb303e2 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -109,10 +109,33 @@ If `trust_remote_code` needs to be true, please append the command with the `--t ### TensorRT-LLM +Uses the `trtllm` backend built into lm-eval (>= 0.4.12), which loads the quantized +checkpoint directly with the TensorRT-LLM LLM API. + ```sh -python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=,checkpoint_dir= --tasks --batch_size +python lm_eval_hf.py --model trtllm \ + --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512 \ + --tasks \ + --batch_size ``` +> **_NOTE:_** Set `max_input_len` and `max_output_len` explicitly. They default to 2048 and +> 512, and prompts longer than `max_input_len` are silently truncated — 5-shot MMLU or +> gsm8k prompts exceed 2048 tokens. `max_seq_len` of the engine is their sum. + +> **_NOTE:_** `tensor_parallel_size` defaults to 1; set it to the number of GPUs the +> checkpoint needs. `pipeline_parallel_size` is also supported. + +> **_NOTE:_** Run this through `lm_eval_hf.py`, not the plain `lm_eval` CLI. lm-eval +> 0.4.12's `trtllm` backend misaligns TensorRT-LLM's `prompt_logprobs` by one position, so +> every loglikelihood task (hellaswag, mmlu, arc, ...) fails with a `KeyError`; +> `lm_eval_hf.py` patches the alignment before the model is built. The patch is removed +> once the fix lands upstream. + +The previous `lm_eval_tensorrt_llm.py` entry point (`--model trt-llm`) is deprecated. It +still works — it prints a warning and forwards the translated command to the above — but it +will be removed in a future release. + ## MMLU [Massive Multitask Language Understanding](https://arxiv.org/abs/2009.03300). A score (0-1, higher is better) will be printed at the end of the benchmark. diff --git a/examples/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py index 3f406b10341..ee7515245a3 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -39,6 +39,7 @@ import contextlib import glob import json +import logging import os import sys import warnings @@ -48,8 +49,10 @@ from lm_eval import utils from packaging.version import Version -if Version(version("lm_eval")) < Version("0.4.10"): - raise ImportError(f"lm_eval_hf.py requires lm-eval >= 0.4.10; found {version('lm_eval')}.") +if Version(version("lm_eval")) < Version("0.4.12"): + # 0.4.12 is the first release shipping lm_eval.models.trtllm_causallms (the `trtllm` + # backend), which replaced this example's own TensorRT-LLM model wrapper. + raise ImportError(f"lm_eval_hf.py requires lm-eval >= 0.4.12; found {version('lm_eval')}.") from lm_eval._cli import HarnessCLI from lm_eval.api.model import T @@ -69,6 +72,8 @@ except ImportError: _ANYMODEL_AVAILABLE = False +logger = logging.getLogger(__name__) + def _anymodel_patcher_context(pretrained, trust_remote_code=False): """Return a deci_x_patcher context if *pretrained* is a Puzzletron checkpoint, else a no-op.""" @@ -292,6 +297,39 @@ def _inject_modelopt_args_into_model_args(args): args.model_args = model_args +def _patch_trtllm_parse_logprobs(): + """Fix the prompt_logprobs off-by-one in lm-eval's `trtllm` backend. + + TensorRT-LLM aligns `prompt_logprobs` to the *next* token: its worker computes them + from `prompt_token_ids[1:] + first_generated_token`, so entry ``i`` is the + distribution that predicted ``tokens[i + 1]`` and contains that token's id. + lm-eval 0.4.12's ``TRTLLM._parse_logprobs`` instead reads ``prompt_logprobs[i] + [tokens[i]]`` and applies its own shift on top, which raises ``KeyError`` on the + first request of every loglikelihood task (hellaswag, mmlu, arc, ...). + + Drop this once the upstream backend indexes the entries correctly. + """ + from lm_eval.models.trtllm_causallms import TRTLLM + + def _parse_logprobs(tokens, outputs, ctxlen): + prompt_logprobs = outputs.outputs[0].prompt_logprobs + continuation_logprobs = 0.0 + is_greedy = True + # Token 0 has no preceding distribution, so it can never be scored. + for i in range(max(ctxlen, 1), len(tokens)): + logprob = prompt_logprobs[i - 1].get(tokens[i]) + if logprob is None: + # TRT-LLM always appends the actual token to the top-k dict. + logger.warning(f"Token {tokens[i]} missing from prompt_logprobs[{i - 1}]") + continue + continuation_logprobs += logprob.logprob + if logprob.rank != 1: + is_greedy = False + return continuation_logprobs, is_greedy + + TRTLLM._parse_logprobs = staticmethod(_parse_logprobs) + + def _enforce_accuracy_gate(output_path, task, lower_bound): """Read lm-eval results at output_path; exit non-zero if /acc < lower_bound.""" # HarnessCLI.execute() returns None, so we recover the result dict from disk. @@ -331,6 +369,9 @@ def _enforce_accuracy_gate(output_path, task, lower_bound): _add_modelopt_args(run_parser) args = cli.parse_args() _inject_modelopt_args_into_model_args(args) + if getattr(args, "model", None) == "trtllm": + # Patch lazily: importing the backend pulls in tensorrt_llm. + _patch_trtllm_parse_logprobs() lower_bound = vars(args).pop("accuracy_lower_bound", None) output_path = getattr(args, "output_path", None) gate_task = None diff --git a/examples/llm_eval/lm_eval_tensorrt_llm.py b/examples/llm_eval/lm_eval_tensorrt_llm.py index d64fdcc02de..61dab651262 100644 --- a/examples/llm_eval/lm_eval_tensorrt_llm.py +++ b/examples/llm_eval/lm_eval_tensorrt_llm.py @@ -13,201 +13,114 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy -import gc -import logging +"""[Deprecated] ModelOpt's TensorRT-LLM backend for lm-evaluation-harness. + +lm-eval >= 0.4.12 ships its own TensorRT-LLM backend +(``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), so this example no +longer maintains one. Migrate to: + + python lm_eval_hf.py --model trtllm \ + --model_args model=,tokenizer=,\ +tensor_parallel_size=,max_batch_size= \ + --tasks --batch_size + +This shim keeps legacy ``--model trt-llm`` commands working by translating them to the +upstream backend and forwarding to ``lm_eval_hf.py``. It will be removed in a future +release. +""" + import os -import signal -import threading -import time -from collections.abc import Iterable -from typing import Any - -import torch -import torch.nn.functional as F -from lm_eval.__main__ import cli_evaluate -from lm_eval.api.registry import register_model -from lm_eval.models.api_models import TemplateAPI -from transformers import BatchEncoding - -from modelopt.deploy.llm import LLM - -logger = logging.getLogger(__name__) - -TokenSequence = list[int] | torch.LongTensor | torch.Tensor | BatchEncoding - - -@register_model("trt-llm") -class TRTLLM(TemplateAPI): - def __init__( - self, - tokenizer: str, - checkpoint_dir: str, - batch_size: int = 1, - **kwargs, - ): - assert isinstance(tokenizer, str) - super().__init__( - tokenizer=tokenizer, - batch_size=int(batch_size), - **kwargs, - ) - - if self.tokenizer.pad_token_id is None: - self.tokenizer.pad_token_id = self.tokenizer.eos_token_id - - assert isinstance(checkpoint_dir, str) - - max_length = kwargs.get("max_length", self._max_gen_toks + 4096) - self.llm = LLM( - checkpoint_dir=checkpoint_dir, - tokenizer=self.tokenizer, - max_batch_size=int(batch_size), - max_seq_len=max_length, - # Loglikelihood tasks request context logits. KV cache prefix reuse would return - # logits only for the recomputed suffix on shared-prefix requests (e.g. hellaswag), - # truncating context_logits and breaking parse_logprobs. Disable it. - enable_kv_cache_reuse=False, - trust_remote_code=bool(kwargs.get("trust_remote_code", False)), - ) - self.max_length = max_length - 1 - logger.info("Loaded TRT-LLM") - - def model_call( - self, - messages: Iterable[list[int]], - *, - generate: bool = True, - gen_kwargs: dict | None = None, - **kwargs, - ): - # !!! Copy: shared dict for each request, need new object !!! - gen_kwargs = copy.deepcopy(gen_kwargs) - - assert isinstance(messages, Iterable), "Expect the messages to be Iterable[list[int]]" - first_element = next(iter(messages)) - assert isinstance(first_element, list) and isinstance(first_element[0], int), ( - "Expect the messages to be Iterable[list[int]]" - ) - - if not generate: - return self.llm.generate_context_logits(prompts=messages) - - llm_kwargs = {} - max_new_tokens = self._max_gen_toks - stop_words = [] - if gen_kwargs: - if "until" in gen_kwargs: - stop_words = gen_kwargs.pop("until") - llm_kwargs["stop_words"] = stop_words - if "temperature" in gen_kwargs: - llm_kwargs["temperature"] = gen_kwargs.pop("temperature") - if "top_p" in gen_kwargs: - llm_kwargs["top_p"] = gen_kwargs.pop("top_p") - if "max_gen_toks" in gen_kwargs: - max_new_tokens = gen_kwargs.pop("max_gen_toks") - - output_texts: list[str] = self.llm.generate_text( - prompts=messages, - max_new_tokens=max_new_tokens, - **llm_kwargs, - ) - - # Manually filter out keyword if not supported by llm. - for i, text in enumerate(output_texts): - for word in stop_words: - word_index = text.find(word) - if word_index >= 0: - text = text[:word_index] - output_texts[i] = text - - return output_texts - - async def amodel_call( - self, - session, - messages: Iterable[list[int]], - *, - generate: bool = True, - cache_keys: list | None = None, - ctxlens: list[int] | None = None, - gen_kwargs: dict | None = None, - **kwargs, - ): - raise NotImplementedError - - def loglikelihood_rolling(self, requests): - raise NotImplementedError - - def _create_payload( - self, - messages: list[list[int]] | list[dict] | list[str] | str, - *, - generate: bool = True, - gen_kwargs: dict | None = None, - seed: int = 1234, - **kwargs, - ) -> dict: - """This method is responsible for creating the json payload that will be sent to the API.""" - raise NotImplementedError - - @staticmethod - def parse_generations(outputs: Any | list[Any], **kwargs) -> list[str]: - """Method used to parse the generations from the (batched) API response.""" - return outputs - - @staticmethod - def parse_logprobs( - outputs: Any | list[Any], - tokens: list[list[int]] | None = None, - ctxlens: list[int] | None = None, - **kwargs, - ) -> list[tuple[float, bool]]: - """Method used to parse the logprobs from the (batched) API response. - - The provided tokens have two parts: The context tokens (length as ctxlens) and the continuation tokens. - The logprobs returned is computed from the continuation tokens. - We return the sum of the logprob of the continuation tokens - [assuming the continuation tokens are the golden output]. - """ - res = [] - - for logits_single_batch, tokens_single_batch, ctxlen_single_batch in zip( - outputs, - tokens, # type: ignore[arg-type] - ctxlens, # type: ignore[arg-type] - ): - logits_single_batch = logits_single_batch.to("cuda") - continuation_logprob = F.log_softmax( - logits_single_batch[(ctxlen_single_batch - 1) : -1], dim=-1 - ) - continuation_tokens = torch.tensor(tokens_single_batch[ctxlen_single_batch:]) - top_tokens = continuation_logprob.argmax(dim=-1).cpu() - - is_greedy = torch.equal(top_tokens, continuation_tokens) - - logprob_sum = ( - continuation_logprob[ - torch.arange(continuation_logprob.size(0)), continuation_tokens - ] - .sum() - .cpu() - ) - - res.append((logprob_sum, is_greedy)) - - return res +import sys +from pathlib import Path +# Legacy --model_args key -> upstream `trtllm` key. `tokenizer` and `trust_remote_code` +# carry over unchanged. +_RENAMED_ARGS = {"checkpoint_dir": "model"} -if __name__ == "__main__": - cli_evaluate() - # Force clean up the LLM instance and void hanging. - gc.collect() +# The old wrapper sized the TensorRT-LLM engine from `max_length` (total sequence +# length), defaulting to max_gen_toks + 4096. Upstream splits that into an input and an +# output budget: max_seq_len = max_input_len + max_output_len. +_DEFAULT_MAX_GEN_TOKS = 256 +_DEFAULT_MAX_INPUT_LEN = 4096 + + +def _parse_model_args(raw: str) -> dict[str, str]: + """Parse lm-eval's comma-separated ``key=value`` --model_args string.""" + args = {} + for field in raw.split(","): + if not field: + continue + key, _, value = field.partition("=") + args[key.strip()] = value.strip() + return args + + +def _translate_model_args(raw: str, batch_size: str | None) -> str: + """Rewrite legacy `trt-llm` --model_args into upstream `trtllm` --model_args.""" + args = _parse_model_args(raw) + args = {_RENAMED_ARGS.get(k, k): v for k, v in args.items()} + + max_gen_toks = int(args.get("max_gen_toks", _DEFAULT_MAX_GEN_TOKS)) + # Upstream defaults to max_input_len=2048, short enough to silently left-truncate + # 5-shot prompts, so always derive it from the legacy sizing instead. + args.setdefault("max_output_len", str(max_gen_toks)) + max_length = args.pop("max_length", None) + args.setdefault( + "max_input_len", + str(int(max_length) - max_gen_toks) if max_length else str(_DEFAULT_MAX_INPUT_LEN), + ) + + # The old wrapper built the engine for --batch_size and used every visible GPU. + if batch_size and batch_size.isdigit(): + args.setdefault("max_batch_size", batch_size) + if "tensor_parallel_size" not in args: + import torch - # Force terminate in case gc.collect() is not enough. - def _terminate(): - time.sleep(10) - os.kill(os.getpid(), signal.SIGTERM) + args.setdefault("tensor_parallel_size", str(max(torch.cuda.device_count(), 1))) - termination_thread = threading.Thread(target=_terminate, daemon=True) - termination_thread.start() + return ",".join(f"{k}={v}" for k, v in args.items()) + + +def _find_option(argv: list[str], name: str) -> tuple[int, str] | None: + """Locate ``--name `` or ``--name=``; return (index, value).""" + for i, arg in enumerate(argv): + if arg == name and i + 1 < len(argv): + return i + 1, argv[i + 1] + if arg.startswith(f"{name}="): + return i, arg[len(name) + 1 :] + return None + + +def _set_option(argv: list[str], name: str, index: int, value: str) -> None: + """Write *value* back at *index*, preserving the ``--name=`` form.""" + argv[index] = f"{name}={value}" if argv[index].startswith(f"{name}=") else value + + +def _translate(argv: list[str]) -> list[str]: + """Rewrite a legacy argv so it targets the upstream `trtllm` backend.""" + out = list(argv) + # --batch_size sized the engine in the old wrapper; read it before rewriting. + found = _find_option(out, "--batch_size") + batch_size = found[1] if found else None + + found = _find_option(out, "--model") + if found and found[1] in ("trt-llm", "trt_llm"): + _set_option(out, "--model", found[0], "trtllm") + + found = _find_option(out, "--model_args") + if found: + _set_option(out, "--model_args", found[0], _translate_model_args(found[1], batch_size)) + return out + + +if __name__ == "__main__": + forwarded = _translate(sys.argv[1:]) + target = Path(__file__).resolve().parent / "lm_eval_hf.py" + print( + "WARNING: examples/llm_eval/lm_eval_tensorrt_llm.py is deprecated and will be " + "removed in a future release.\n" + " lm-eval >= 0.4.12 provides the `trtllm` backend; forwarding to:\n" + f" python {target.name} {' '.join(forwarded)}", + file=sys.stderr, + ) + os.execv(sys.executable, [sys.executable, str(target), *forwarded]) diff --git a/examples/llm_eval/requirements.txt b/examples/llm_eval/requirements.txt index 2762c838c6a..5fc4a03c3ca 100644 --- a/examples/llm_eval/requirements.txt +++ b/examples/llm_eval/requirements.txt @@ -1,5 +1,5 @@ fire>=0.5.0 -lm_eval[api,ifeval]>=0.4.10 +lm_eval[api,ifeval]>=0.4.12,<0.5 peft>=0.5.0 rwkv>=0.7.3 torchvision diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py index 61934f2b141..d65f497e320 100644 --- a/tests/examples/llm_eval/test_llm_eval.py +++ b/tests/examples/llm_eval/test_llm_eval.py @@ -48,8 +48,9 @@ def test_lm_eval_hf(tmp_path): def test_qwen3_eval_fp8(tmp_path): # Bump max_position_embeddings: TRT-LLM serve rejects prompts longer than max_seq_len. # The default (32) is shorter than even simple MMLU prompts, and 2048 is shorter than - # 5-shot gsm8k prompts (~3.9k tokens). The eval LLM caps max_seq_len at max_gen_toks + 4096, - # so 8192 leaves headroom for the longest prompts we evaluate. + # 5-shot gsm8k prompts (~3.9k tokens). huggingface_example.sh sizes the lm_eval engine + # at max_input_len 4096 + max_output_len, so 8192 leaves headroom for the longest + # prompts we evaluate. model_dir = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, max_position_embeddings=8192) try: run_hf_ptq_command( From dcedd37b4b1bec7df2d7929d5fa99498dc86537e Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 4 Aug 2026 23:06:14 +0000 Subject: [PATCH 2/7] Move the trtllm logprob fix out of lm_eval_hf.py, drop the legacy shim Review feedback: - lm_eval_hf.py is the HF entry point and should not carry a TensorRT-LLM patch. The prompt_logprobs override moves to a new lm_eval_trtllm.py, whose only content is that override plus lm-eval's CLI; lm_eval_hf.py is back to HF-only. - Delete lm_eval_tensorrt_llm.py outright instead of keeping a translating shim. README and CHANGELOG give the replacement command. - huggingface_example.sh hardcoded max_input_len=4096. parser.sh already echoed BUILD_MAX_INPUT_LEN but never parsed or defaulted it, so wire up --input (default 4096) and use it. - Use MODEL_ABS_PATH for the tokenizer: the block runs after `pushd ../llm_eval/`, so a relative --model would not resolve (the mmlu block already does this). Quote --model_args (SC2086). - LM_EVAL_TP is now overridable, and the README documents the TensorRT-LLM tuning the deleted wrapper applied that the upstream backend does not expose (expert parallelism, KV-cache fraction). The override also now fails loudly: it checks prompt_logprobs covers every prompt token and raises when a token is missing, instead of skipping the term and silently inflating the reported accuracy. Re-verified on nvidia/Qwen3.5-122B-A10B-NVFP4 (4x B300, TRT-LLM 1.3.0rc23), tp=4: hellaswag acc 0.7188 / acc_norm 0.7812, unchanged from the previous commit and from the deleted implementation. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../skills/deployment/references/trtllm.md | 6 +- CHANGELOG.rst | 2 +- .../hf_ptq/scripts/huggingface_example.sh | 14 +- examples/hf_ptq/scripts/parser.sh | 7 +- examples/llm_eval/README.md | 21 +-- examples/llm_eval/lm_eval_hf.py | 42 +----- examples/llm_eval/lm_eval_tensorrt_llm.py | 126 ------------------ examples/llm_eval/lm_eval_trtllm.py | 84 ++++++++++++ 8 files changed, 116 insertions(+), 186 deletions(-) delete mode 100644 examples/llm_eval/lm_eval_tensorrt_llm.py create mode 100644 examples/llm_eval/lm_eval_trtllm.py diff --git a/.agents/skills/deployment/references/trtllm.md b/.agents/skills/deployment/references/trtllm.md index 59027774fb1..74895b30dab 100644 --- a/.agents/skills/deployment/references/trtllm.md +++ b/.agents/skills/deployment/references/trtllm.md @@ -70,10 +70,12 @@ If you encounter a legacy checkpoint (no `hf_quant_config.json`, has `rank*.safe ## Evaluation with TRT-LLM Runs lm-evaluation-harness benchmarks through lm-eval's built-in `trtllm` backend -(requires `lm_eval>=0.4.12`; ModelOpt's own `lm_eval_tensorrt_llm.py` is deprecated). +(requires `lm_eval>=0.4.12`; ModelOpt's own `lm_eval_tensorrt_llm.py` has been removed). +`lm_eval_trtllm.py` is a thin wrapper that corrects the backend's `prompt_logprobs` +alignment — without it every loglikelihood task raises `KeyError`. ```bash -python examples/llm_eval/lm_eval_hf.py \ +python examples/llm_eval/lm_eval_trtllm.py \ --model trtllm \ --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512 \ --tasks gsm8k,mmlu \ diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 30129f067a1..62ae44f146d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,7 +19,7 @@ Changelog **Deprecations** -- Deprecate ``examples/llm_eval/lm_eval_tensorrt_llm.py`` (the ``trt-llm`` model). lm-evaluation-harness 0.4.12 ships its own TensorRT-LLM backend (``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), which also implements ``loglikelihood_rolling`` and pipeline parallelism, so the example uses that instead and pins ``lm_eval>=0.4.12,<0.5``. The old entry point still runs: it warns, translates the legacy arguments (``checkpoint_dir`` → ``model``, ``max_length`` → ``max_input_len``/``max_output_len``, ``--batch_size`` → ``max_batch_size``, tensor parallelism over all visible GPUs) and forwards to ``lm_eval_hf.py --model trtllm``. When calling the new backend directly, set ``tensor_parallel_size`` and ``max_input_len`` explicitly — they default to 1 and 2048, and longer prompts are silently truncated. ``lm_eval_hf.py`` also patches that backend's ``_parse_logprobs``, which is off by one against TensorRT-LLM's next-token-aligned ``prompt_logprobs`` and otherwise raises ``KeyError`` on every loglikelihood task; run TensorRT-LLM evaluations through ``lm_eval_hf.py`` rather than the plain ``lm_eval`` CLI until that is fixed upstream. +- Remove ``examples/llm_eval/lm_eval_tensorrt_llm.py`` (the ``trt-llm`` model). lm-evaluation-harness 0.4.12 ships its own TensorRT-LLM backend (``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), which also implements ``loglikelihood_rolling`` and pipeline parallelism, so the example uses that instead and pins ``lm_eval>=0.4.12,<0.5``. Replace ``python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=,checkpoint_dir=`` with ``python lm_eval_trtllm.py --model trtllm --model_args model=,tokenizer=``, and set ``tensor_parallel_size`` and ``max_input_len`` explicitly — they default to 1 and 2048, and longer prompts are silently truncated. ``lm_eval_trtllm.py`` is a thin entry point that overrides the backend's ``_parse_logprobs``, which is off by one against TensorRT-LLM's next-token-aligned ``prompt_logprobs`` and otherwise raises ``KeyError`` on every loglikelihood task; use it rather than the plain ``lm_eval`` CLI until that is fixed upstream. ``examples/hf_ptq/scripts/huggingface_example.sh`` gains ``--input`` (``BUILD_MAX_INPUT_LEN``, default 4096) to size the evaluation engine's context, and honours a preset ``LM_EVAL_TP`` to override the tensor-parallel size. **Bug Fixes** diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index fadb97d0070..44363cb012e 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -298,15 +298,17 @@ if [[ $TASKS =~ "lm_eval" ]]; then pip install -r requirements.txt # lm-eval's `trtllm` backend defaults to 1 GPU; shard over every visible one instead. - LM_EVAL_TP=$(python -c "import torch; print(max(torch.cuda.device_count(), 1))") + # Override LM_EVAL_TP to lower it -- TRT-LLM enables expert parallelism at higher TP, + # which fails in DeepEP kernels for MoE checkpoints on some GPUs (e.g. SM 12.0). + LM_EVAL_TP=${LM_EVAL_TP:-$(python -c "import torch; print(max(torch.cuda.device_count(), 1))")} - echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE tp $LM_EVAL_TP" + echo "Using the following config: max input $BUILD_MAX_INPUT_LEN max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE tp $LM_EVAL_TP" - # max_input_len defaults to 2048, which silently truncates 5-shot prompts; size the - # engine for a 4096-token context plus the requested generation length. - python lm_eval_hf.py \ + # max_input_len defaults to 2048, which silently truncates 5-shot prompts, so pass it + # explicitly; the engine's max_seq_len is max_input_len + max_output_len. + python lm_eval_trtllm.py \ --model trtllm \ - --model_args model=$SAVE_PATH,tokenizer=$MODEL_PATH,tensor_parallel_size=$LM_EVAL_TP,max_batch_size=$BUILD_MAX_BATCH_SIZE,max_gen_toks=$BUILD_MAX_OUTPUT_LEN,max_input_len=4096,max_output_len=$BUILD_MAX_OUTPUT_LEN \ + --model_args "model=$SAVE_PATH,tokenizer=$MODEL_ABS_PATH,tensor_parallel_size=$LM_EVAL_TP,max_batch_size=$BUILD_MAX_BATCH_SIZE,max_gen_toks=$BUILD_MAX_OUTPUT_LEN,max_input_len=$BUILD_MAX_INPUT_LEN,max_output_len=$BUILD_MAX_OUTPUT_LEN" \ --tasks $LM_EVAL_TASKS \ --batch_size $BUILD_MAX_BATCH_SIZE $lm_eval_flags | tee $LM_EVAL_RESULT diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index 03ed3a57631..af54832dfd9 100644 --- a/examples/hf_ptq/scripts/parser.sh +++ b/examples/hf_ptq/scripts/parser.sh @@ -41,7 +41,7 @@ parse_options() { CALIB_WITH_IMAGES=false # Parse command-line options - ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,input:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") eval set -- "$ARGS" while true; do @@ -56,6 +56,7 @@ parse_options() { --awq_block_size ) AWQ_BLOCK_SIZE="$2"; shift 2;; --calib ) CALIB_SIZE="$2"; shift 2;; --calib_batch_size ) CALIB_BATCH_SIZE="$2"; shift 2;; + --input ) BUILD_MAX_INPUT_LEN="$2"; shift 2;; --output ) BUILD_MAX_OUTPUT_LEN="$2"; shift 2;; --batch ) BUILD_MAX_BATCH_SIZE="$2"; shift 2;; --tasks ) TASKS="$2"; shift 2;; @@ -90,6 +91,7 @@ parse_options() { DEFAULT_CALIB_SIZE=512 DEFAULT_CALIB_SEQ=512 DEFAULT_CALIB_BATCH_SIZE=0 + DEFAULT_BUILD_MAX_INPUT_LEN=4096 DEFAULT_BUILD_MAX_OUTPUT_LEN=1024 DEFAULT_BUILD_MAX_BATCH_SIZE=2 @@ -102,6 +104,9 @@ parse_options() { if [ -z "$CALIB_BATCH_SIZE" ]; then CALIB_BATCH_SIZE=$DEFAULT_CALIB_BATCH_SIZE fi + if [ -z "$BUILD_MAX_INPUT_LEN" ]; then + BUILD_MAX_INPUT_LEN=$DEFAULT_BUILD_MAX_INPUT_LEN + fi if [ -z "$BUILD_MAX_OUTPUT_LEN" ]; then BUILD_MAX_OUTPUT_LEN=$DEFAULT_BUILD_MAX_OUTPUT_LEN fi diff --git a/examples/llm_eval/README.md b/examples/llm_eval/README.md index af0dfb303e2..1ab143d23fa 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -113,7 +113,7 @@ Uses the `trtllm` backend built into lm-eval (>= 0.4.12), which loads the quanti checkpoint directly with the TensorRT-LLM LLM API. ```sh -python lm_eval_hf.py --model trtllm \ +python lm_eval_trtllm.py --model trtllm \ --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512 \ --tasks \ --batch_size @@ -126,15 +126,18 @@ python lm_eval_hf.py --model trtllm \ > **_NOTE:_** `tensor_parallel_size` defaults to 1; set it to the number of GPUs the > checkpoint needs. `pipeline_parallel_size` is also supported. -> **_NOTE:_** Run this through `lm_eval_hf.py`, not the plain `lm_eval` CLI. lm-eval -> 0.4.12's `trtllm` backend misaligns TensorRT-LLM's `prompt_logprobs` by one position, so -> every loglikelihood task (hellaswag, mmlu, arc, ...) fails with a `KeyError`; -> `lm_eval_hf.py` patches the alignment before the model is built. The patch is removed -> once the fix lands upstream. +> **_NOTE:_** Use `lm_eval_trtllm.py` rather than the plain `lm_eval` CLI. lm-eval 0.4.12's +> `trtllm` backend misaligns TensorRT-LLM's `prompt_logprobs` by one position, so every +> loglikelihood task (hellaswag, mmlu, arc, ...) fails with a `KeyError`; +> `lm_eval_trtllm.py` overrides the alignment. It goes away once the fix lands upstream. -The previous `lm_eval_tensorrt_llm.py` entry point (`--model trt-llm`) is deprecated. It -still works — it prints a warning and forwards the translated command to the above — but it -will be removed in a future release. +> **_NOTE:_** The backend forwards only a fixed set of arguments to TensorRT-LLM, so the +> tuning the old `lm_eval_tensorrt_llm.py` applied is not reachable: expert parallelism is +> left at the TensorRT-LLM default (MoE checkpoints can fail in DeepEP kernels on some +> GPUs, e.g. SM 12.0) and the KV cache uses 90% of free GPU memory rather than 70%. Lower +> `tensor_parallel_size` if you hit either. + +`lm_eval_tensorrt_llm.py` (`--model trt-llm`) has been removed; use the command above. ## MMLU diff --git a/examples/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py index ee7515245a3..a96a730c16a 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -39,7 +39,6 @@ import contextlib import glob import json -import logging import os import sys import warnings @@ -50,8 +49,7 @@ from packaging.version import Version if Version(version("lm_eval")) < Version("0.4.12"): - # 0.4.12 is the first release shipping lm_eval.models.trtllm_causallms (the `trtllm` - # backend), which replaced this example's own TensorRT-LLM model wrapper. + # Matches the floor pinned in requirements.txt. raise ImportError(f"lm_eval_hf.py requires lm-eval >= 0.4.12; found {version('lm_eval')}.") from lm_eval._cli import HarnessCLI @@ -72,8 +70,6 @@ except ImportError: _ANYMODEL_AVAILABLE = False -logger = logging.getLogger(__name__) - def _anymodel_patcher_context(pretrained, trust_remote_code=False): """Return a deci_x_patcher context if *pretrained* is a Puzzletron checkpoint, else a no-op.""" @@ -297,39 +293,6 @@ def _inject_modelopt_args_into_model_args(args): args.model_args = model_args -def _patch_trtllm_parse_logprobs(): - """Fix the prompt_logprobs off-by-one in lm-eval's `trtllm` backend. - - TensorRT-LLM aligns `prompt_logprobs` to the *next* token: its worker computes them - from `prompt_token_ids[1:] + first_generated_token`, so entry ``i`` is the - distribution that predicted ``tokens[i + 1]`` and contains that token's id. - lm-eval 0.4.12's ``TRTLLM._parse_logprobs`` instead reads ``prompt_logprobs[i] - [tokens[i]]`` and applies its own shift on top, which raises ``KeyError`` on the - first request of every loglikelihood task (hellaswag, mmlu, arc, ...). - - Drop this once the upstream backend indexes the entries correctly. - """ - from lm_eval.models.trtllm_causallms import TRTLLM - - def _parse_logprobs(tokens, outputs, ctxlen): - prompt_logprobs = outputs.outputs[0].prompt_logprobs - continuation_logprobs = 0.0 - is_greedy = True - # Token 0 has no preceding distribution, so it can never be scored. - for i in range(max(ctxlen, 1), len(tokens)): - logprob = prompt_logprobs[i - 1].get(tokens[i]) - if logprob is None: - # TRT-LLM always appends the actual token to the top-k dict. - logger.warning(f"Token {tokens[i]} missing from prompt_logprobs[{i - 1}]") - continue - continuation_logprobs += logprob.logprob - if logprob.rank != 1: - is_greedy = False - return continuation_logprobs, is_greedy - - TRTLLM._parse_logprobs = staticmethod(_parse_logprobs) - - def _enforce_accuracy_gate(output_path, task, lower_bound): """Read lm-eval results at output_path; exit non-zero if /acc < lower_bound.""" # HarnessCLI.execute() returns None, so we recover the result dict from disk. @@ -369,9 +332,6 @@ def _enforce_accuracy_gate(output_path, task, lower_bound): _add_modelopt_args(run_parser) args = cli.parse_args() _inject_modelopt_args_into_model_args(args) - if getattr(args, "model", None) == "trtllm": - # Patch lazily: importing the backend pulls in tensorrt_llm. - _patch_trtllm_parse_logprobs() lower_bound = vars(args).pop("accuracy_lower_bound", None) output_path = getattr(args, "output_path", None) gate_task = None diff --git a/examples/llm_eval/lm_eval_tensorrt_llm.py b/examples/llm_eval/lm_eval_tensorrt_llm.py deleted file mode 100644 index 61dab651262..00000000000 --- a/examples/llm_eval/lm_eval_tensorrt_llm.py +++ /dev/null @@ -1,126 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""[Deprecated] ModelOpt's TensorRT-LLM backend for lm-evaluation-harness. - -lm-eval >= 0.4.12 ships its own TensorRT-LLM backend -(``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), so this example no -longer maintains one. Migrate to: - - python lm_eval_hf.py --model trtllm \ - --model_args model=,tokenizer=,\ -tensor_parallel_size=,max_batch_size= \ - --tasks --batch_size - -This shim keeps legacy ``--model trt-llm`` commands working by translating them to the -upstream backend and forwarding to ``lm_eval_hf.py``. It will be removed in a future -release. -""" - -import os -import sys -from pathlib import Path - -# Legacy --model_args key -> upstream `trtllm` key. `tokenizer` and `trust_remote_code` -# carry over unchanged. -_RENAMED_ARGS = {"checkpoint_dir": "model"} - -# The old wrapper sized the TensorRT-LLM engine from `max_length` (total sequence -# length), defaulting to max_gen_toks + 4096. Upstream splits that into an input and an -# output budget: max_seq_len = max_input_len + max_output_len. -_DEFAULT_MAX_GEN_TOKS = 256 -_DEFAULT_MAX_INPUT_LEN = 4096 - - -def _parse_model_args(raw: str) -> dict[str, str]: - """Parse lm-eval's comma-separated ``key=value`` --model_args string.""" - args = {} - for field in raw.split(","): - if not field: - continue - key, _, value = field.partition("=") - args[key.strip()] = value.strip() - return args - - -def _translate_model_args(raw: str, batch_size: str | None) -> str: - """Rewrite legacy `trt-llm` --model_args into upstream `trtllm` --model_args.""" - args = _parse_model_args(raw) - args = {_RENAMED_ARGS.get(k, k): v for k, v in args.items()} - - max_gen_toks = int(args.get("max_gen_toks", _DEFAULT_MAX_GEN_TOKS)) - # Upstream defaults to max_input_len=2048, short enough to silently left-truncate - # 5-shot prompts, so always derive it from the legacy sizing instead. - args.setdefault("max_output_len", str(max_gen_toks)) - max_length = args.pop("max_length", None) - args.setdefault( - "max_input_len", - str(int(max_length) - max_gen_toks) if max_length else str(_DEFAULT_MAX_INPUT_LEN), - ) - - # The old wrapper built the engine for --batch_size and used every visible GPU. - if batch_size and batch_size.isdigit(): - args.setdefault("max_batch_size", batch_size) - if "tensor_parallel_size" not in args: - import torch - - args.setdefault("tensor_parallel_size", str(max(torch.cuda.device_count(), 1))) - - return ",".join(f"{k}={v}" for k, v in args.items()) - - -def _find_option(argv: list[str], name: str) -> tuple[int, str] | None: - """Locate ``--name `` or ``--name=``; return (index, value).""" - for i, arg in enumerate(argv): - if arg == name and i + 1 < len(argv): - return i + 1, argv[i + 1] - if arg.startswith(f"{name}="): - return i, arg[len(name) + 1 :] - return None - - -def _set_option(argv: list[str], name: str, index: int, value: str) -> None: - """Write *value* back at *index*, preserving the ``--name=`` form.""" - argv[index] = f"{name}={value}" if argv[index].startswith(f"{name}=") else value - - -def _translate(argv: list[str]) -> list[str]: - """Rewrite a legacy argv so it targets the upstream `trtllm` backend.""" - out = list(argv) - # --batch_size sized the engine in the old wrapper; read it before rewriting. - found = _find_option(out, "--batch_size") - batch_size = found[1] if found else None - - found = _find_option(out, "--model") - if found and found[1] in ("trt-llm", "trt_llm"): - _set_option(out, "--model", found[0], "trtllm") - - found = _find_option(out, "--model_args") - if found: - _set_option(out, "--model_args", found[0], _translate_model_args(found[1], batch_size)) - return out - - -if __name__ == "__main__": - forwarded = _translate(sys.argv[1:]) - target = Path(__file__).resolve().parent / "lm_eval_hf.py" - print( - "WARNING: examples/llm_eval/lm_eval_tensorrt_llm.py is deprecated and will be " - "removed in a future release.\n" - " lm-eval >= 0.4.12 provides the `trtllm` backend; forwarding to:\n" - f" python {target.name} {' '.join(forwarded)}", - file=sys.stderr, - ) - os.execv(sys.executable, [sys.executable, str(target), *forwarded]) diff --git a/examples/llm_eval/lm_eval_trtllm.py b/examples/llm_eval/lm_eval_trtllm.py new file mode 100644 index 00000000000..68a3ccafdb5 --- /dev/null +++ b/examples/llm_eval/lm_eval_trtllm.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run lm-evaluation-harness against a TensorRT-LLM checkpoint. + +Entry point around lm-eval's built-in ``trtllm`` backend +(``lm_eval.models.trtllm_causallms``, new in 0.4.12). It exists only to correct that +backend's ``prompt_logprobs`` handling -- everything else is upstream. Drop this file and +call ``lm_eval`` directly once the fix lands upstream. + + python lm_eval_trtllm.py --model trtllm \ + --model_args model=,tokenizer=,\ +tensor_parallel_size=,max_batch_size=,max_input_len=4096 \ + --tasks --batch_size +""" + +from importlib.metadata import version + +from lm_eval.__main__ import cli_evaluate +from packaging.version import Version + +if Version(version("lm_eval")) < Version("0.4.12"): + # 0.4.12 is the first release shipping lm_eval.models.trtllm_causallms. + raise ImportError(f"lm_eval_trtllm.py requires lm-eval >= 0.4.12; found {version('lm_eval')}.") + +from lm_eval.models.trtllm_causallms import TRTLLM + + +def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, bool]: + """Sum the continuation logprobs of one request, correcting upstream's alignment. + + TensorRT-LLM aligns ``prompt_logprobs`` to the *next* token: its worker computes them + from ``prompt_token_ids[1:] + first_generated_token`` (``executor/base_worker.py``), so + entry ``i`` is the distribution that predicted ``tokens[i + 1]`` and always contains + that token's id -- either in the top-k or appended by ``_topk_logprobs``. + + lm-eval 0.4.12's ``TRTLLM._parse_logprobs`` instead reads + ``prompt_logprobs[i][tokens[i]]`` and applies its own shift on top, which raises + ``KeyError`` on the first request of every loglikelihood task (hellaswag, mmlu, arc). + """ + prompt_logprobs = outputs.outputs[0].prompt_logprobs + # Scoring tokens[ctxlen:] reads entries ctxlen-1 .. len(tokens)-2; a shorter list means + # the engine saw a different prompt than we asked about, which would shift every index. + if len(prompt_logprobs) < len(tokens) - 1: + raise RuntimeError( + f"prompt_logprobs has {len(prompt_logprobs)} entries for {len(tokens)} tokens; " + "the engine scored a different prompt than was requested." + ) + + continuation_logprobs = 0.0 + is_greedy = True + # Token 0 has no preceding distribution, so it can never be scored. + for i in range(max(ctxlen, 1), len(tokens)): + logprob = prompt_logprobs[i - 1].get(tokens[i]) + if logprob is None: + # Dropping the term instead would silently inflate the reported accuracy. + raise RuntimeError( + f"tokens[{i}] is missing from prompt_logprobs[{i - 1}]; the returned " + "logprobs are misaligned with the requested tokens." + ) + continuation_logprobs += logprob.logprob + if logprob.rank != 1: + is_greedy = False + + return continuation_logprobs, is_greedy + + +TRTLLM._parse_logprobs = staticmethod(_parse_logprobs) + + +if __name__ == "__main__": + cli_evaluate() From 622b97c26456f9520136e866260077bdeea485a6 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 5 Aug 2026 16:48:01 +0000 Subject: [PATCH 3/7] Unit-test the trtllm logprob override and guard against upstream drift Review feedback: the corrected _parse_logprobs is the only new logic in this PR and had no test of its own -- the GPU numbers exercise it end to end but cannot pin the alignment. tests/examples/llm_eval/test_lm_eval_trtllm.py stubs the TensorRT-LLM response object (no GPU, no tensorrt_llm install) and covers the i-1 alignment, the rank != 1 -> is_greedy rule, the ctxlen=0 edge, and both RuntimeError paths. Mutation-checked: dropping the -1 shift is caught by 5/5 cases, ignoring ctxlen by 4/5. The last test is a tripwire. lm_eval_trtllm keeps the upstream implementation in _UPSTREAM_PARSE_LOGPROBS and asserts it still raises KeyError on a correctly-shaped response. A future 0.4.x that fixes the alignment would otherwise be silently re-broken by this override, since requirements.txt only pins >=0.4.12,<0.5; instead that test fails and says to delete the file. Also raise a clear error if TRTLLM._parse_logprobs disappears, rather than overriding an attribute that no longer exists. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- examples/llm_eval/lm_eval_trtllm.py | 11 ++ .../examples/llm_eval/test_lm_eval_trtllm.py | 121 ++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 tests/examples/llm_eval/test_lm_eval_trtllm.py diff --git a/examples/llm_eval/lm_eval_trtllm.py b/examples/llm_eval/lm_eval_trtllm.py index 68a3ccafdb5..18a930b6dfb 100644 --- a/examples/llm_eval/lm_eval_trtllm.py +++ b/examples/llm_eval/lm_eval_trtllm.py @@ -77,6 +77,17 @@ def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, boo return continuation_logprobs, is_greedy +if not hasattr(TRTLLM, "_parse_logprobs"): + raise RuntimeError( + "lm_eval.models.trtllm_causallms.TRTLLM has no _parse_logprobs to override; the " + f"backend changed shape in lm-eval {version('lm_eval')}. Recheck whether this file " + "is still needed." + ) + +# Kept so the unit tests can assert the upstream implementation is still the broken one. +# When that assertion starts failing, upstream has fixed the alignment and this whole file +# should be deleted in favour of calling `lm_eval` directly. +_UPSTREAM_PARSE_LOGPROBS = TRTLLM._parse_logprobs TRTLLM._parse_logprobs = staticmethod(_parse_logprobs) diff --git a/tests/examples/llm_eval/test_lm_eval_trtllm.py b/tests/examples/llm_eval/test_lm_eval_trtllm.py new file mode 100644 index 00000000000..4fabc2f83a3 --- /dev/null +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for ``examples/llm_eval/lm_eval_trtllm.py``. + +The module lives next to the example script (not inside the ``modelopt`` package), +so we add ``examples/llm_eval/`` to ``sys.path`` before importing it. No GPU and no +TensorRT-LLM install is needed: ``_parse_logprobs`` is pure Python over a response +object, which these tests stub out. +""" + +import sys +from dataclasses import dataclass +from pathlib import Path + +import pytest + +pytest.importorskip("lm_eval", reason="lm_eval is an examples/llm_eval requirement") + +_LLM_EVAL_DIR = Path(__file__).resolve().parents[3] / "examples" / "llm_eval" +if str(_LLM_EVAL_DIR) not in sys.path: + sys.path.insert(0, str(_LLM_EVAL_DIR)) + +import lm_eval_trtllm + +# TensorRT-LLM aligns prompt_logprobs to the *next* token, so entry i holds the +# distribution that predicted tokens[i + 1]. These fixtures mirror that layout. +TOKENS = [10, 11, 12, 13, 14] +CTXLEN = 3 # continuation is tokens[3:] == [13, 14] + + +@dataclass +class _Logprob: + """Stand-in for ``tensorrt_llm.executor.result.Logprob``.""" + + logprob: float + rank: int = 1 + + +class _Outputs: + """Stand-in for a TensorRT-LLM ``RequestOutput``.""" + + def __init__(self, prompt_logprobs): + self.outputs = [type("_Completion", (), {"prompt_logprobs": prompt_logprobs})()] + + +def _prompt_logprobs(continuation_ranks=(1, 1)): + """One dict per prompt token; entry i contains tokens[i + 1], as TRT-LLM returns.""" + return [ + {TOKENS[1]: _Logprob(-2.0)}, # predicts tokens[1] -- context, not scored + {TOKENS[2]: _Logprob(-3.0)}, # predicts tokens[2] -- context, not scored + {TOKENS[3]: _Logprob(-0.5, continuation_ranks[0])}, # predicts tokens[3] -- scored + {TOKENS[4]: _Logprob(-1.25, continuation_ranks[1])}, # predicts tokens[4] -- scored + {999: _Logprob(-9.0)}, # predicts the first generated token -- unused + ] + + +def test_sums_only_continuation_tokens(): + """Only tokens[ctxlen:] contribute, read from the entry one position earlier.""" + logprob, is_greedy = lm_eval_trtllm._parse_logprobs( + tokens=TOKENS, outputs=_Outputs(_prompt_logprobs()), ctxlen=CTXLEN + ) + assert logprob == pytest.approx(-0.5 + -1.25) + assert is_greedy is True + + +def test_is_greedy_false_when_a_continuation_token_is_not_rank_one(): + logprob, is_greedy = lm_eval_trtllm._parse_logprobs( + tokens=TOKENS, + outputs=_Outputs(_prompt_logprobs(continuation_ranks=(1, 2))), + ctxlen=CTXLEN, + ) + assert logprob == pytest.approx(-0.5 + -1.25) + assert is_greedy is False + + +def test_ctxlen_zero_skips_the_unscorable_first_token(): + """tokens[0] has no preceding distribution, so scoring starts at tokens[1].""" + logprob, _ = lm_eval_trtllm._parse_logprobs( + tokens=TOKENS, outputs=_Outputs(_prompt_logprobs()), ctxlen=0 + ) + assert logprob == pytest.approx(-2.0 + -3.0 + -0.5 + -1.25) + + +def test_raises_when_prompt_logprobs_is_too_short(): + """A short list means the engine scored a different prompt; every index would shift.""" + with pytest.raises(RuntimeError, match="entries for 5 tokens"): + lm_eval_trtllm._parse_logprobs( + tokens=TOKENS, outputs=_Outputs(_prompt_logprobs()[:3]), ctxlen=CTXLEN + ) + + +def test_raises_when_a_continuation_token_is_missing(): + """Dropping the term instead would silently inflate the reported accuracy.""" + entries = _prompt_logprobs() + entries[2] = {777: _Logprob(-0.5)} # tokens[3] absent from the entry that predicts it + with pytest.raises(RuntimeError, match=r"tokens\[3\] is missing"): + lm_eval_trtllm._parse_logprobs(tokens=TOKENS, outputs=_Outputs(entries), ctxlen=CTXLEN) + + +def test_upstream_is_still_misaligned(): + """Tripwire: when this fails, upstream fixed the bug and this file can be deleted. + + lm-eval's own ``_parse_logprobs`` reads ``prompt_logprobs[i][tokens[i]]``, but entry i + holds ``tokens[i + 1]``, so it raises ``KeyError`` on the very first prompt token. + """ + with pytest.raises(KeyError): + lm_eval_trtllm._UPSTREAM_PARSE_LOGPROBS( + tokens=TOKENS, outputs=_Outputs(_prompt_logprobs()), ctxlen=CTXLEN + ) From 196c3260630f2aabf4f8d373467100c517843c79 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 5 Aug 2026 17:18:27 +0000 Subject: [PATCH 4/7] Pin --trust_remote_code propagation for the trtllm entry point Review question: lm_eval_trtllm.py calls lm-eval's CLI directly and does no argument rewriting, while lm_eval_hf.py goes out of its way to move --trust_remote_code into model_args. Nothing confirmed the flag that huggingface_example.sh passes actually reaches the trtllm tokenizer. It does, on lm-eval 0.4.12: _cli/run.py:355 builds the config via EvaluatorConfig.from_cli, which calls _configure() -> _set_trust_remote_code() (config/evaluate_config.py:274-276, 420-435). That sets datasets.config.HF_DATASETS_TRUST_REMOTE_CODE and injects model_args["trust_remote_code"]=True, which TRTLLM.__init__ takes as a named parameter and forwards to both AutoTokenizer.from_pretrained and the TensorRT-LLM LLM kwargs (trtllm_causallms.py:48, 110, 148). lm_eval_hf.py's manual handling is therefore redundant on 0.4.12 -- it predates lm-eval doing this itself. Left alone here; it is harmless (it populates model_args and then nulls the namespace flag, so the built-in path no-ops) and out of scope for this PR. Three tests pin the contract we now depend on: the flag is injected when set, not injected when unset, and trust_remote_code is still a parameter of TRTLLM.__init__ so the injected key lands somewhere real. No GPU required. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../examples/llm_eval/test_lm_eval_trtllm.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/examples/llm_eval/test_lm_eval_trtllm.py b/tests/examples/llm_eval/test_lm_eval_trtllm.py index 4fabc2f83a3..086abfb3743 100644 --- a/tests/examples/llm_eval/test_lm_eval_trtllm.py +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -109,6 +109,49 @@ def test_raises_when_a_continuation_token_is_missing(): lm_eval_trtllm._parse_logprobs(tokens=TOKENS, outputs=_Outputs(entries), ctxlen=CTXLEN) +def test_trust_remote_code_reaches_the_backend(monkeypatch): + """`--trust_remote_code` must land in model_args, since we call lm-eval's CLI directly. + + `huggingface_example.sh` passes the flag, and unlike `lm_eval_hf.py` this entry point + does no rewriting of its own -- it relies on lm-eval doing it. On the live path that is + `_cli/run.py:355` -> `EvaluatorConfig.from_cli` -> `_configure()` -> + `_set_trust_remote_code()`. + """ + import datasets + from lm_eval.config.evaluate_config import EvaluatorConfig + + monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False) + + cfg = EvaluatorConfig( + model="trtllm", model_args={"model": "/ckpt", "tokenizer": "/tok"}, trust_remote_code=True + ) + cfg._set_trust_remote_code() + + assert cfg.model_args["trust_remote_code"] is True + assert datasets.config.HF_DATASETS_TRUST_REMOTE_CODE is True + + +def test_trust_remote_code_not_injected_when_unset(monkeypatch): + import datasets + from lm_eval.config.evaluate_config import EvaluatorConfig + + monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False) + + cfg = EvaluatorConfig(model="trtllm", model_args={"model": "/ckpt"}, trust_remote_code=False) + cfg._set_trust_remote_code() + + assert "trust_remote_code" not in cfg.model_args + + +def test_trtllm_backend_accepts_trust_remote_code(): + """The key lm-eval injects has to be a parameter the backend actually takes.""" + import inspect + + from lm_eval.models.trtllm_causallms import TRTLLM + + assert "trust_remote_code" in inspect.signature(TRTLLM.__init__).parameters + + def test_upstream_is_still_misaligned(): """Tripwire: when this fails, upstream fixed the bug and this file can be deleted. From d5f868cf1e28fd45a8ff13d5b40c4a48f2bb281b Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 6 Aug 2026 22:59:41 +0000 Subject: [PATCH 5/7] Skip the trtllm tests on the backend module, not just lm_eval Review nit: lm_eval.models.trtllm_causallms currently guards its tensorrt_llm import, but if that ever becomes eager, importorskip("lm_eval") would let collection proceed and then error. Skip on the module actually imported. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- tests/examples/llm_eval/test_lm_eval_trtllm.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/examples/llm_eval/test_lm_eval_trtllm.py b/tests/examples/llm_eval/test_lm_eval_trtllm.py index 086abfb3743..8f8c2754dba 100644 --- a/tests/examples/llm_eval/test_lm_eval_trtllm.py +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -26,7 +26,11 @@ import pytest -pytest.importorskip("lm_eval", reason="lm_eval is an examples/llm_eval requirement") +# Skip on the backend module, not just `lm_eval`: it currently guards its `tensorrt_llm` +# import, but if that ever becomes eager this should skip rather than error. +pytest.importorskip( + "lm_eval.models.trtllm_causallms", reason="lm_eval is an examples/llm_eval requirement" +) _LLM_EVAL_DIR = Path(__file__).resolve().parents[3] / "examples" / "llm_eval" if str(_LLM_EVAL_DIR) not in sys.path: From 6a54747c3fadbfca5a635993154dc762c499a02e Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 7 Aug 2026 18:51:23 +0000 Subject: [PATCH 6/7] Require TensorRT-LLM >= 1.3.0rc11 and fix the tests on datasets 4.x Two P1 review findings, both reproduced locally before fixing. 1. The i-1 prompt_logprobs mapping this file corrects is only valid from TensorRT-LLM 1.3.0rc11. Confirmed against the tags: v1.2.0 calls compute_logprobs(prompt_logprobs, None, context_logits, None, None) with no prompt token ids, so _topk_logprobs skips the "append the requested token" branch and prompt_logprobs=1 keeps only the argmax; a non-greedy continuation token is simply absent. hf_ptq/README.md recommends release:1.2.0, so a user following the docs would hit this. _check_trtllm_version() raises with an actionable message. It runs from _parse_logprobs rather than at startup, so generative-only runs -- which never touch this path -- keep working on older releases, and __main__ warns up front so an unusable container is visible before the model loads. Both READMEs and the CHANGELOG now state the requirement. 2. The trust_remote_code tests failed on datasets 4.x, which dropped HF_DATASETS_TRUST_REMOTE_CODE; monkeypatch.setattr then raised AttributeError before exercising anything. Reproduced by deleting the attribute (7 passed, 2 failed, matching the report) and fixed with raising=False -- lm-eval creates the attribute on assignment, and monkeypatch still removes it on teardown. Seven new cases cover the version guard, including that it fires from _parse_logprobs and not only from __main__. A missing tensorrt_llm is treated as "nothing to check" so the suite still needs no TensorRT-LLM install; that branch is unreachable in a real run, where the backend refuses to build a model without it. 17 tests pass under datasets 3.3.2 and under a simulated 4.x. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- examples/hf_ptq/README.md | 5 +++ examples/llm_eval/README.md | 6 +++ examples/llm_eval/lm_eval_trtllm.py | 41 +++++++++++++++++++ .../examples/llm_eval/test_lm_eval_trtllm.py | 39 +++++++++++++++++- 5 files changed, 90 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 62ae44f146d..df21947167e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,7 +19,7 @@ Changelog **Deprecations** -- Remove ``examples/llm_eval/lm_eval_tensorrt_llm.py`` (the ``trt-llm`` model). lm-evaluation-harness 0.4.12 ships its own TensorRT-LLM backend (``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), which also implements ``loglikelihood_rolling`` and pipeline parallelism, so the example uses that instead and pins ``lm_eval>=0.4.12,<0.5``. Replace ``python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=,checkpoint_dir=`` with ``python lm_eval_trtllm.py --model trtllm --model_args model=,tokenizer=``, and set ``tensor_parallel_size`` and ``max_input_len`` explicitly — they default to 1 and 2048, and longer prompts are silently truncated. ``lm_eval_trtllm.py`` is a thin entry point that overrides the backend's ``_parse_logprobs``, which is off by one against TensorRT-LLM's next-token-aligned ``prompt_logprobs`` and otherwise raises ``KeyError`` on every loglikelihood task; use it rather than the plain ``lm_eval`` CLI until that is fixed upstream. ``examples/hf_ptq/scripts/huggingface_example.sh`` gains ``--input`` (``BUILD_MAX_INPUT_LEN``, default 4096) to size the evaluation engine's context, and honours a preset ``LM_EVAL_TP`` to override the tensor-parallel size. +- Remove ``examples/llm_eval/lm_eval_tensorrt_llm.py`` (the ``trt-llm`` model). lm-evaluation-harness 0.4.12 ships its own TensorRT-LLM backend (``lm_eval.models.trtllm_causallms``, registered as ``trtllm``), which also implements ``loglikelihood_rolling`` and pipeline parallelism, so the example uses that instead and pins ``lm_eval>=0.4.12,<0.5``. Replace ``python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=,checkpoint_dir=`` with ``python lm_eval_trtllm.py --model trtllm --model_args model=,tokenizer=``, and set ``tensor_parallel_size`` and ``max_input_len`` explicitly — they default to 1 and 2048, and longer prompts are silently truncated. ``lm_eval_trtllm.py`` is a thin entry point that overrides the backend's ``_parse_logprobs``, which is off by one against TensorRT-LLM's next-token-aligned ``prompt_logprobs`` and otherwise raises ``KeyError`` on every loglikelihood task; use it rather than the plain ``lm_eval`` CLI until that is fixed upstream. Loglikelihood tasks additionally require **TensorRT-LLM >= 1.3.0rc11**, the release that started passing prompt token ids into ``compute_logprobs`` so the scored token is present in every entry; older releases return only the top-1 token per position and the run aborts with an explicit error instead of a confusing lookup failure. Generative tasks are unaffected. ``examples/hf_ptq/scripts/huggingface_example.sh`` gains ``--input`` (``BUILD_MAX_INPUT_LEN``, default 4096) to size the evaluation engine's context, and honours a preset ``LM_EVAL_TP`` to override the tensor-parallel size. **Bug Fixes** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index e80926bea91..5d752bc5f11 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -31,6 +31,11 @@ This section focuses on Post-training quantization, a technique that reduces mod For Hugging Face models, please use the TensorRT-LLM docker image (e.g., `nvcr.io/nvidia/tensorrt-llm/release:1.2.0`). Visit our [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more information. +> **NOTE:** `--tasks lm_eval` needs **TensorRT-LLM >= 1.3.0rc11** for loglikelihood +> benchmarks (mmlu, hellaswag, arc, ...). Earlier releases return only the top-1 token per +> prompt position, so a continuation token's logprob cannot be recovered. Generative +> benchmarks (gsm8k, ifeval) work on older releases. + Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies. ### Local Installation diff --git a/examples/llm_eval/README.md b/examples/llm_eval/README.md index 1ab143d23fa..f4c53f49d81 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -119,6 +119,12 @@ python lm_eval_trtllm.py --model trtllm \ --batch_size ``` +> **_NOTE:_** Loglikelihood tasks (mmlu, hellaswag, arc, ...) need **TensorRT-LLM >= +> 1.3.0rc11**, which is when the engine started returning the requested token in every +> `prompt_logprobs` entry. Earlier releases return only the top-1 token per position, so a +> continuation token's logprob cannot be recovered and the run aborts with a clear error. +> Generative tasks (gsm8k, ifeval) are unaffected. + > **_NOTE:_** Set `max_input_len` and `max_output_len` explicitly. They default to 2048 and > 512, and prompts longer than `max_input_len` are silently truncated — 5-shot MMLU or > gsm8k prompts exceed 2048 tokens. `max_seq_len` of the engine is their sum. diff --git a/examples/llm_eval/lm_eval_trtllm.py b/examples/llm_eval/lm_eval_trtllm.py index 18a930b6dfb..b9a4b2a92da 100644 --- a/examples/llm_eval/lm_eval_trtllm.py +++ b/examples/llm_eval/lm_eval_trtllm.py @@ -26,6 +26,7 @@ --tasks --batch_size """ +import sys from importlib.metadata import version from lm_eval.__main__ import cli_evaluate @@ -37,6 +38,32 @@ from lm_eval.models.trtllm_causallms import TRTLLM +# TensorRT-LLM only started passing the prompt token ids into `compute_logprobs` in +# 1.3.0rc11 (`executor/base_worker.py`), which is what makes the requested token always +# present in each `prompt_logprobs` entry. On 1.2.0 and earlier, `prompt_logprobs=1` keeps +# only the top-1 token, so a non-greedy continuation token is simply absent and no correct +# continuation logprob can be recovered -- by this file or by lm-eval's own version. +_MIN_TRTLLM_VERSION = "1.3.0rc11" +_trtllm_version_checked = False + + +def _check_trtllm_version() -> None: + """Raise if TensorRT-LLM predates the `prompt_logprobs` layout scored below.""" + try: + import tensorrt_llm + except ImportError: + # Nothing to check, and unreachable in a real run: the backend refuses to build a + # model without tensorrt_llm long before any logprob is scored. + return + + if Version(tensorrt_llm.__version__) < Version(_MIN_TRTLLM_VERSION): + raise RuntimeError( + f"Loglikelihood tasks need TensorRT-LLM >= {_MIN_TRTLLM_VERSION}; found " + f"{tensorrt_llm.__version__}. Earlier releases return only the top-1 token per " + "prompt position, so the continuation token's logprob is unavailable. Use a " + "newer TensorRT-LLM container, or restrict the run to generative tasks." + ) + def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, bool]: """Sum the continuation logprobs of one request, correcting upstream's alignment. @@ -50,6 +77,13 @@ def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, boo ``prompt_logprobs[i][tokens[i]]`` and applies its own shift on top, which raises ``KeyError`` on the first request of every loglikelihood task (hellaswag, mmlu, arc). """ + global _trtllm_version_checked + if not _trtllm_version_checked: + # Checked here rather than at startup so generative-only runs, which never reach + # this path, still work on older TensorRT-LLM releases. + _check_trtllm_version() + _trtllm_version_checked = True + prompt_logprobs = outputs.outputs[0].prompt_logprobs # Scoring tokens[ctxlen:] reads entries ctxlen-1 .. len(tokens)-2; a shorter list means # the engine saw a different prompt than we asked about, which would shift every index. @@ -92,4 +126,11 @@ def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, boo if __name__ == "__main__": + # Warn up front so an unusable container is obvious before the model loads, but do not + # abort: generative tasks are unaffected by the old prompt_logprobs layout. + try: + _check_trtllm_version() + except RuntimeError as e: + print(f"WARNING: {e}", file=sys.stderr) + cli_evaluate() diff --git a/tests/examples/llm_eval/test_lm_eval_trtllm.py b/tests/examples/llm_eval/test_lm_eval_trtllm.py index 8f8c2754dba..19f0978e16f 100644 --- a/tests/examples/llm_eval/test_lm_eval_trtllm.py +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -21,6 +21,7 @@ """ import sys +import types from dataclasses import dataclass from pathlib import Path @@ -124,7 +125,9 @@ def test_trust_remote_code_reaches_the_backend(monkeypatch): import datasets from lm_eval.config.evaluate_config import EvaluatorConfig - monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False) + # raising=False: datasets 4.x dropped this attribute, and lm-eval creates it on + # assignment. monkeypatch still removes it again on teardown. + monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False, raising=False) cfg = EvaluatorConfig( model="trtllm", model_args={"model": "/ckpt", "tokenizer": "/tok"}, trust_remote_code=True @@ -139,7 +142,9 @@ def test_trust_remote_code_not_injected_when_unset(monkeypatch): import datasets from lm_eval.config.evaluate_config import EvaluatorConfig - monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False) + # raising=False: datasets 4.x dropped this attribute, and lm-eval creates it on + # assignment. monkeypatch still removes it again on teardown. + monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False, raising=False) cfg = EvaluatorConfig(model="trtllm", model_args={"model": "/ckpt"}, trust_remote_code=False) cfg._set_trust_remote_code() @@ -156,6 +161,36 @@ def test_trtllm_backend_accepts_trust_remote_code(): assert "trust_remote_code" in inspect.signature(TRTLLM.__init__).parameters +def _fake_trtllm(monkeypatch, version): + module = types.ModuleType("tensorrt_llm") + module.__version__ = version + monkeypatch.setitem(sys.modules, "tensorrt_llm", module) + + +@pytest.mark.parametrize("version", ["1.1.0rc2", "1.2.0", "1.3.0rc10"]) +def test_rejects_trtllm_without_the_prompt_token_ids_fix(monkeypatch, version): + """<1.3.0rc11 returns only the top-1 token per position, so continuations are absent.""" + _fake_trtllm(monkeypatch, version) + with pytest.raises(RuntimeError, match=r"1\.3\.0rc11"): + lm_eval_trtllm._check_trtllm_version() + + +@pytest.mark.parametrize("version", ["1.3.0rc11", "1.3.0rc23", "1.3.0", "1.4.0"]) +def test_accepts_trtllm_with_the_prompt_token_ids_fix(monkeypatch, version): + _fake_trtllm(monkeypatch, version) + lm_eval_trtllm._check_trtllm_version() + + +def test_version_is_checked_before_scoring(monkeypatch): + """The guard has to fire from _parse_logprobs, not only from __main__.""" + _fake_trtllm(monkeypatch, "1.2.0") + monkeypatch.setattr(lm_eval_trtllm, "_trtllm_version_checked", False) + with pytest.raises(RuntimeError, match=r"1\.3\.0rc11"): + lm_eval_trtllm._parse_logprobs( + tokens=TOKENS, outputs=_Outputs(_prompt_logprobs()), ctxlen=CTXLEN + ) + + def test_upstream_is_still_misaligned(): """Tripwire: when this fails, upstream fixed the bug and this file can be deleted. From 482a84ed1b3ec93b07d7a5d5cf8f9a4899940f89 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 7 Aug 2026 21:37:20 +0000 Subject: [PATCH 7/7] Document the TensorRT-LLM version requirement only in llm_eval Drop the note added to examples/hf_ptq/README.md; the requirement is stated in examples/llm_eval/README.md, which is where the evaluation command lives, and in the CHANGELOG. hf_ptq/README.md is now untouched by this PR. Users on an older container still get the actionable runtime error from lm_eval_trtllm.py rather than a confusing missing-token failure. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- examples/hf_ptq/README.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 5d752bc5f11..e80926bea91 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -31,11 +31,6 @@ This section focuses on Post-training quantization, a technique that reduces mod For Hugging Face models, please use the TensorRT-LLM docker image (e.g., `nvcr.io/nvidia/tensorrt-llm/release:1.2.0`). Visit our [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more information. -> **NOTE:** `--tasks lm_eval` needs **TensorRT-LLM >= 1.3.0rc11** for loglikelihood -> benchmarks (mmlu, hellaswag, arc, ...). Earlier releases return only the top-1 token per -> prompt position, so a continuation token's logprob cannot be recovered. Generative -> benchmarks (gsm8k, ifeval) work on older releases. - Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies. ### Local Installation