diff --git a/.agents/skills/deployment/references/trtllm.md b/.agents/skills/deployment/references/trtllm.md index 865c11e7eee..74895b30dab 100644 --- a/.agents/skills/deployment/references/trtllm.md +++ b/.agents/skills/deployment/references/trtllm.md @@ -69,14 +69,22 @@ 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` 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_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 \ + --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..df21947167e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,8 @@ 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. 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** 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..44363cb012e 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -297,11 +297,18 @@ 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" - - 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 \ + # lm-eval's `trtllm` backend defaults to 1 GPU; shard over every visible one instead. + # 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 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, 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_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 f680235b651..f4c53f49d81 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -109,10 +109,42 @@ 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_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 ``` +> **_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. + +> **_NOTE:_** `tensor_parallel_size` defaults to 1; set it to the number of GPUs the +> checkpoint needs. `pipeline_parallel_size` is also supported. + +> **_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. + +> **_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 [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..a96a730c16a 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -48,8 +48,9 @@ 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"): + # 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 from lm_eval.api.model import T 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 d64fdcc02de..00000000000 --- a/examples/llm_eval/lm_eval_tensorrt_llm.py +++ /dev/null @@ -1,213 +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. - -import copy -import gc -import logging -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 - - -if __name__ == "__main__": - cli_evaluate() - # Force clean up the LLM instance and void hanging. - gc.collect() - - # Force terminate in case gc.collect() is not enough. - def _terminate(): - time.sleep(10) - os.kill(os.getpid(), signal.SIGTERM) - - termination_thread = threading.Thread(target=_terminate, daemon=True) - termination_thread.start() diff --git a/examples/llm_eval/lm_eval_trtllm.py b/examples/llm_eval/lm_eval_trtllm.py new file mode 100644 index 00000000000..b9a4b2a92da --- /dev/null +++ b/examples/llm_eval/lm_eval_trtllm.py @@ -0,0 +1,136 @@ +# 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 +""" + +import sys +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 + +# 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. + + 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). + """ + 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. + 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 + + +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) + + +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/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( 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..19f0978e16f --- /dev/null +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -0,0 +1,203 @@ +# 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 +import types +from dataclasses import dataclass +from pathlib import Path + +import pytest + +# 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: + 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_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 + + # 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 + ) + 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 + + # 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() + + 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 _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. + + 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 + )