Use lm-eval 0.4.12's built-in trtllm backend, deprecate lm_eval_tensorrt_llm.py - #2066
Conversation
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughTensorRT-LLM evaluation now uses lm-evaluation-harness’s built-in ChangesTensorRT-LLM evaluation migration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant EvaluationScript
participant lm_eval_trtllm.py
participant lm-evaluation-harness
participant TRTLLMBackend
EvaluationScript->>lm_eval_trtllm.py: pass model, tokenizer, and runtime limits
lm_eval_trtllm.py->>TRTLLMBackend: install corrected prompt-logprob parser
lm_eval_trtllm.py->>lm-evaluation-harness: invoke cli_evaluate()
lm-evaluation-harness->>TRTLLMBackend: run trtllm evaluation
TRTLLMBackend-->>lm_eval_trtllm.py: return aligned log probabilities
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🧹 Nitpick comments (2)
examples/llm_eval/lm_eval_hf.py (1)
314-328: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winClear
is_greedywhen a token is missing.The fallback branch skips the token. The score then omits a negative log-probability, and
is_greedycan still be reported asTrue. A skipped token means the rank is unknown, so the greedy claim is not supported. Setis_greedy = Falsein that branch.♻️ Proposed change
if logprob is None: # TRT-LLM always appends the actual token to the top-k dict. logger.warning(f"Token at position {i} missing from prompt_logprobs[{i - 1}]") + is_greedy = False continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_eval/lm_eval_hf.py` around lines 314 - 328, Update the missing-token branch in _parse_logprobs to set is_greedy = False before continuing, while preserving the existing warning and score-skipping behavior.examples/llm_eval/lm_eval_tensorrt_llm.py (1)
76-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the reason for the local
torchimport.The coding guidelines allow a local import only for a justified circular dependency, optional dependency, or unusually heavy import, and they require a brief explanatory comment. Add the comment here.
As per coding guidelines: "use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment".
♻️ Proposed change
if "tensor_parallel_size" not in args: + # Local import: torch is heavy and only needed to count visible GPUs. import torch args.setdefault("tensor_parallel_size", str(max(torch.cuda.device_count(), 1)))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_eval/lm_eval_tensorrt_llm.py` around lines 76 - 79, Add a brief explanatory comment immediately before the local torch import in the tensor_parallel_size setup, stating the applicable justification for keeping this import local. Leave the import and surrounding device-count logic unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/scripts/huggingface_example.sh`:
- Around line 307-309: Update the lm_eval_hf.py invocation to use the existing
MODEL_ABS_PATH variable for tokenizer and quote the complete --model_args value,
preserving all existing model argument settings while preventing path resolution
and word-splitting issues.
In `@examples/llm_eval/lm_eval_hf.py`:
- Around line 321-324: The logger.warning call in the token processing logic
includes the token ID from tokens[i] in the message, but the review requests
logging only the token position instead. Update the warning message to remove
the token value and report only the position index i - 1 where the logprob
lookup failed.
---
Nitpick comments:
In `@examples/llm_eval/lm_eval_hf.py`:
- Around line 314-328: Update the missing-token branch in _parse_logprobs to set
is_greedy = False before continuing, while preserving the existing warning and
score-skipping behavior.
In `@examples/llm_eval/lm_eval_tensorrt_llm.py`:
- Around line 76-79: Add a brief explanatory comment immediately before the
local torch import in the tensor_parallel_size setup, stating the applicable
justification for keeping this import local. Leave the import and surrounding
device-count logic unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4fe4fe98-cc39-409f-ba89-f42ab2adb6f4
📒 Files selected for processing (8)
.agents/skills/deployment/references/trtllm.mdCHANGELOG.rstexamples/hf_ptq/scripts/huggingface_example.shexamples/llm_eval/README.mdexamples/llm_eval/lm_eval_hf.pyexamples/llm_eval/lm_eval_tensorrt_llm.pyexamples/llm_eval/requirements.txttests/examples/llm_eval/test_llm_eval.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2066 +/- ##
==========================================
- Coverage 65.51% 65.37% -0.14%
==========================================
Files 521 522 +1
Lines 59812 63843 +4031
==========================================
+ Hits 39185 41740 +2555
- Misses 20627 22103 +1476
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/claude review |
kevalmorabia97
left a comment
kevalmorabia97
left a comment
There was a problem hiding this comment.
LGTM. Please review unresolved coderabbit / claude contents
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # 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 |
There was a problem hiding this comment.
[IMPORTANT Algorithm] The patched loop indexes prompt_logprobs[i - 1] with i taken from lm-eval's full tokens list, but nothing checks that prompt_logprobs actually has one entry per prompt token. Two things in this PR make that assumption fragile:
- The PR description itself notes upstream's
maybe_truncateleft-truncates prompts longer thanmax_input_len. If the engine sees a truncated prompt whiletokensis untruncated, every index shifts by the number of dropped tokens — soprompt_logprobs[i-1]returns the distribution for the wrong position. When the shifted dict happens to containtokens[i](very common: repeated tokens in few-shot prompts) you get a plausible-but-wrong logprob with no signal at all; when it doesn't, you hit thecontinuebelow and silently drop the term fromcontinuation_logprobs. Both inflate the score rather than failing. - The deleted wrapper set
enable_kv_cache_reuse=Falseprecisely because prefix-block reuse returns logits only for the recomputed suffix on shared-prefix requests (hellaswag is exactly that shape). The upstream backend has no equivalent knob and TRT-LLM defaultsenable_block_reuse=True, so that invariant is now unenforced and depends entirely on upstream/TRT-LLM behavior.
The continue fallback turns any such misalignment into a quietly optimistic accuracy number, which is the worst outcome for an eval script. A cheap loud check makes the whole patch self-verifying:
def _parse_logprobs(tokens, outputs, ctxlen):
prompt_logprobs = outputs.outputs[0].prompt_logprobs
# prompt_logprobs is next-token aligned and must cover every prompt token; a
# shorter list means the engine saw a different prompt (e.g. max_input_len
# truncation or KV-cache prefix reuse), which would silently 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 prompt was likely truncated — raise max_input_len."
)
...
logprob = prompt_logprobs[i - 1].get(tokens[i])
if logprob is None:
raise RuntimeError(
f"tokens[{i}] missing from prompt_logprobs[{i - 1}]; prompt_logprobs "
"are misaligned with the request tokens."
)If you'd rather keep the tolerant path, at minimum make it fail the run (or return -inf) instead of skipping the token — a warning in a multi-thousand-line lm-eval log will not be noticed, and the reported acc will look normal.
There was a problem hiding this comment.
Took the hardening, pushing back on the two premises.
Taken (dcedd37b4): the override now checks len(prompt_logprobs) >= len(tokens) - 1 and raises on a missing token instead of continue. You are right that silently dropping a term inflates the score, which is the worst failure mode for an eval script.
Premise 1 — maybe_truncate shifting the indices — does not apply to this path. maybe_truncate is only called in generate_until. _loglikelihood_tokens truncates the request itself:
inp = (context_enc + continuation_enc)[-(self.max_length) :]and passes that same inp as tokens to _parse_logprobs. The engine and the caller therefore see the same token list by construction; they cannot disagree about truncation.
Premise 2 — enable_block_reuse — is factually wrong. The upstream backend does have the equivalent knob and already sets it, at trtllm_causallms.py:139-142:
# If reuse kv cache, no logprobs can be computed.
kv_cache_config = KvCacheConfig(
enable_block_reuse=False,
)So the invariant the deleted wrapper enforced via enable_kv_cache_reuse=False is still enforced, by upstream, for the same stated reason.
The length check goes in anyway as cheap insurance, but neither named mechanism can produce the misalignment.
| # 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 \ |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This now shards over every visible GPU while dropping two settings the removed wrapper set deliberately, and the combination is what makes the loss visible.
modelopt/deploy/llm/generate.py forced moe_expert_parallel_size=1 with an explicit reason — "Force ep=1 to avoid TRT-LLM DeepEP kernel failures on unsupported GPUs (e.g. Blackwell SM 12.0)" — and set enable_attention_dp=True whenever the config had expert counts. The upstream trtllm backend exposes neither (the PR's "Feature parity notes" confirms this), so TRT-LLM's defaults apply. Setting tensor_parallel_size to the full device count is exactly the configuration where expert parallelism kicks in, so an MoE checkpoint evaluated on a multi-GPU SM 12.0 box can now fail in DeepEP kernels on a command that worked before this PR. free_gpu_memory_fraction likewise moves from the wrapper's 0.7 (with a capped kv_cache.max_tokens) to TRT-LLM's 0.9, raising OOM risk on smaller GPUs.
The CI coverage (test_qwen3_eval_fp8, tiny dense Qwen3) cannot catch either.
Two options, either is fine:
- Make the TP default overridable so a user hitting this has an escape hatch that doesn't require editing the script, e.g.
LM_EVAL_TP=${LM_EVAL_TP:-$(python -c "import torch; print(max(torch.cuda.device_count(), 1))")}. - Or note the lost settings where users will see them (the
examples/llm_eval/README.mdTensorRT-LLM section already has three NOTE blocks; a fourth covering "MoE on unsupported GPUs / OOM → lowertensor_parallel_size" would land the PR-description caveat in the docs).
Right now the caveat lives only in the PR body, so it disappears at merge.
There was a problem hiding this comment.
Did both, in dcedd37b4.
Escape hatch: LM_EVAL_TP=${LM_EVAL_TP:-$(python -c "import torch; print(max(torch.cuda.device_count(), 1))")}, with a comment naming the DeepEP/SM 12.0 failure.
Docs: examples/llm_eval/README.md gains a fourth NOTE covering the tuning the upstream backend does not expose — expert parallelism left at the TRT-LLM default, and the KV cache at 90% of free memory rather than the wrapper's 70% — telling users to lower tensor_parallel_size if they hit either. Agreed the caveat should not have lived only in the PR body.
| # 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 \ |
There was a problem hiding this comment.
[IMPORTANT Compatibility] tokenizer=$MODEL_PATH should be $MODEL_ABS_PATH, and the whole --model_args value should be quoted.
MODEL_ABS_PATH is computed at line 266 (immediately above this block, realpath when MODEL_PATH is a directory) specifically because the script pushd ../llm_eval/ before invoking — a relative --model <dir> resolves against the wrong cwd. The MMLU block a few lines below correctly passes --model_path $MODEL_ABS_PATH; this new line uses the raw $MODEL_PATH, so ./scripts/huggingface_example.sh --model ./my_model ... fails to find the tokenizer. The pre-change command had the same bug, so this is inherited rather than introduced — but the fix is one word and the surrounding code already establishes the convention.
Separately, unquoted $SAVE_PATH/$MODEL_PATH word-split on paths containing spaces, and with set -o pipefail (set earlier in this script) that becomes a confusing arg-parse failure rather than a clear error.
| --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=4096,max_output_len=$BUILD_MAX_OUTPUT_LEN" \ |
There was a problem hiding this comment.
Both fixed in dcedd37b4 (same change as the CodeRabbit thread on this line): tokenizer=$MODEL_ABS_PATH and the whole --model_args value quoted.
Correct that it was inherited — the deleted lm_eval_tensorrt_llm.py invocation had the same $MODEL_PATH bug — but the line was being rewritten anyway, so it is fixed rather than carried forward.
max_input_len is no longer hardcoded either; it comes from the new --input / BUILD_MAX_INPUT_LEN option, which parser.sh echoed but never actually parsed or defaulted.
| 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)) |
There was a problem hiding this comment.
[SUGGESTION] --model_args is translated unconditionally, even when --model was not the legacy trt-llm/trt_llm. So python lm_eval_tensorrt_llm.py --model hf --model_args pretrained=foo gets tensor_parallel_size / max_input_len / max_output_len / max_batch_size injected into HF model args, which then fails inside HFLM.__init__ rather than surfacing the actual mistake ("this script only ever supported trt-llm"). Gating the rewrite on the branch above keeps the shim's behavior confined to the commands it is emulating:
def _translate(argv: list[str]) -> list[str]:
"""Rewrite a legacy argv so it targets the upstream `trtllm` backend."""
out = list(argv)
found = _find_option(out, "--model")
if not found or found[1] not in ("trt-llm", "trt_llm"):
# Not a legacy trt-llm command; forward untouched.
return out
_set_option(out, "--model", found[0], "trtllm")
# --batch_size sized the engine in the old wrapper; read it before rewriting.
found_bs = _find_option(out, "--batch_size")
batch_size = found_bs[1] if found_bs else None
found = _find_option(out, "--model_args")
if found:
_set_option(out, "--model_args", found[0], _translate_model_args(found[1], batch_size))
return outAlso worth a one-line comment at line 74: --batch_size auto is a supported lm-eval value, and the isdigit() guard silently leaves max_batch_size at the upstream default in that case (rather than the engine size the old wrapper would have used) — intentional, but not obvious from the code.
There was a problem hiding this comment.
Moot — lm_eval_tensorrt_llm.py was deleted outright in dcedd37b4 rather than kept as a translating shim, so there is no _translate left to gate.
| 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), | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] max_gen_toks stays in the translated --model_args (it is not in _RENAMED_ARGS and is never popped), so it is forwarded to upstream TRTLLM.__init__ in addition to being consumed here for sizing. Per the PR description, TRTLLM.__init__ takes **kwargs and forwards only a fixed set to the TRT-LLM LLM API — so this is harmless today (silently dropped, or honored as lm-eval's own generation cap), but it does mean an unrecognized-looking key ends up in the forwarded command that gets echoed in the deprecation warning. A short comment saying max_gen_toks is deliberately kept (it is a valid lm-eval LM arg, not just legacy sizing input) would save the next reader the round trip.
Two smaller edge cases in the same block:
int(args.get("max_gen_toks", ...))andint(max_length)will raise a bareValueErroron a non-integer legacy value, from a script whose only job is to be forgiving of legacy input. Atry/exceptre-raising with the offendingkey=valuewould be kinder.max_length - max_gen_tokscan go non-positive (max_length=128,max_gen_toks=256), producingmax_input_len=-128, which then fails deep inside TRT-LLM engine construction. Amax(..., 1)or an explicit error here localizes it.
There was a problem hiding this comment.
Moot — lm_eval_tensorrt_llm.py was deleted outright in dcedd37b4, so the max_gen_toks passthrough and the int() / negative-max_input_len edge cases go with it. max_input_len is now supplied directly by the caller (--input / BUILD_MAX_INPUT_LEN).
| _add_modelopt_args(run_parser) | ||
| args = cli.parse_args() | ||
| _inject_modelopt_args_into_model_args(args) | ||
| if getattr(args, "model", None) == "trtllm": |
There was a problem hiding this comment.
[SUGGESTION] The patch is gated on args.model == "trtllm" exactly, but _inject_modelopt_args_into_model_args just above (line 271) already had to enumerate ("hf", "hf-auto", "huggingface", "hf-multimodal") for the same reason — lm-eval registers backends under multiple aliases. If upstream ever registers an alias for this backend (or a user passes it with different casing), the patch silently does not apply and the run dies with the very KeyError this code exists to prevent — with no hint that the guard was the cause.
Since a missing patch is a hard failure rather than a degradation, consider keying off the resolved class instead of the CLI string, or at least normalizing:
if str(getattr(args, "model", "")).strip().lower() in ("trtllm", "trt-llm"):Non-blocking — the exact string is correct for 0.4.12.
There was a problem hiding this comment.
Moot — the patch is no longer in lm_eval_hf.py, so there is no args.model string guard to normalize. lm_eval_trtllm.py installs the override unconditionally at import, which also removes the failure mode you describe (a guard that silently does not fire).
There was a problem hiding this comment.
Claude review — 0 CRITICAL, 3 IMPORTANT, 3 SUGGESTION
Scope: trigger comment was a bare /claude review, so full review. 8 files changed; I opened the 5 non-doc ones (examples/llm_eval/lm_eval_hf.py, examples/llm_eval/lm_eval_tensorrt_llm.py, examples/hf_ptq/scripts/huggingface_example.sh, examples/llm_eval/requirements.txt, tests/examples/llm_eval/test_llm_eval.py) plus the removed wrapper's replacement, modelopt/deploy/llm/generate.py, for the parity comparison. CHANGELOG.rst, examples/llm_eval/README.md, and .agents/skills/deployment/references/trtllm.md read as accurate and consistent with the code; I did not review them line by line.
Note: the branch is 3 commits behind main, so a two-dot diff shows unrelated examples/vllm_serve/ and tests/gpu_vllm/ files as reverted. I used the gh pr view file list as authoritative and ignored those.
The core change is well-executed. The _parse_logprobs off-by-one diagnosis is correct and specific — prompt_logprobs[i-1][tokens[i]] is the right indexing for TRT-LLM's next-token-aligned entries, and the identical hellaswag numbers across tp=1/2/4 and against the removed implementation is strong evidence, not just a smoke test. The shim's --name=<value> / --name <value> handling and os.execv forwarding are clean, and the removal of ~200 lines of custom backend is a real maintenance win.
Most impactful findings
1. The _parse_logprobs patch has no alignment guard, and its fallback fails silently upward (examples/llm_eval/lm_eval_hf.py:318-324) — IMPORTANT Algorithm
The patch assumes len(prompt_logprobs) == len(tokens) - 1, but this same PR documents that upstream left-truncates prompts past max_input_len via maybe_truncate. Under truncation every index shifts, and because few-shot prompts repeat tokens heavily, the shifted dict will often contain tokens[i] — yielding a plausible-but-wrong logprob with zero signal. When it doesn't, the continue drops the term from the sum. Both directions inflate accuracy rather than failing.
Compounding it: the deleted wrapper set enable_kv_cache_reuse=False precisely because prefix-block reuse returns logits only for the recomputed suffix on shared-prefix requests (hellaswag's exact shape). The upstream backend has no such knob and TRT-LLM defaults enable_block_reuse=True, so that invariant is now unenforced. A two-line length check plus raising instead of continue makes the patch self-verifying.
2. Silently lost engine settings on the multi-GPU MoE path (examples/hf_ptq/scripts/huggingface_example.sh:300-309) — IMPORTANT Compatibility
modelopt/deploy/llm/generate.py forced moe_expert_parallel_size=1 with an explicit reason ("avoid TRT-LLM DeepEP kernel failures on unsupported GPUs (e.g. Blackwell SM 12.0)") and set enable_attention_dp=True for expert configs. Neither is expressible upstream. Defaulting tensor_parallel_size to the full device count is exactly where EP engages, so an MoE checkpoint on a multi-GPU SM 12.0 box can now fail on a command that worked before. free_gpu_memory_fraction also moves 0.7 → 0.9. The PR body documents all of this; the repo does not — make LM_EVAL_TP overridable via env and/or add the caveat to the README's TensorRT-LLM NOTE blocks so it survives the merge.
3. tokenizer=$MODEL_PATH should be $MODEL_ABS_PATH (huggingface_example.sh:309) — IMPORTANT Compatibility
MODEL_ABS_PATH is computed 40 lines above for exactly this reason (the script pushd ../llm_eval/ first), and the MMLU block below uses it correctly. Pre-existing bug carried into the new command, but it's a one-word fix in a line you're already touching. --model_args should also be quoted.
Suggestions (non-blocking)
_translaterewrites--model_argseven when--modelisn't the legacytrt-llm, injecting TRT-LLM keys into an HF command.int()on legacymax_gen_toks/max_lengthraises bareValueError;max_length - max_gen_tokscan go negative →max_input_len=-128.- The
args.model == "trtllm"guard is exact-match, while_inject_modelopt_args_into_model_argsright above it had to enumerate 4 aliases for the same backend-naming reason. A missed patch is a hardKeyError, not a degradation.
Minor, no comment posted: bumping the lm_eval_hf.py guard to >= 0.4.12 hard-fails users on 0.4.10/0.4.11 who only use --model hf and never touch trtllm. requirements.txt is the contract for examples, so this is defensible — just noting the blast radius is wider than the trtllm path.
Risk assessment
Low-to-moderate, and confined to examples/. No modelopt/ source, no mode registration, no modelopt_state schema, no export path — so nothing here can affect checkpoint round-tripping or the library API. Backward compatibility for legacy commands is genuinely preserved and was validated end to end. The residual risk is concentrated in the eval-correctness path: finding #1 can produce quietly-wrong accuracy numbers under conditions this PR makes reachable (truncation, KV reuse), and CI's tiny dense Qwen3 cannot detect either that or finding #2. Adding the length/lookup assertions would move this to low risk outright, since the failure mode becomes loud instead of silent.
| # 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 \ |
There was a problem hiding this comment.
do we have BUILD_MAX_INPUT_LEN here as well?
There was a problem hiding this comment.
Good catch — and it was more broken than it looked. parser.sh:163 already echoed BUILD_MAX_INPUT_LEN, but it was never in the getopt list, never parsed and never defaulted, so it printed empty on every run.
dcedd37b4 wires it up: --input -> BUILD_MAX_INPUT_LEN, DEFAULT_BUILD_MAX_INPUT_LEN=4096, and the lm_eval command now passes max_input_len=$BUILD_MAX_INPUT_LEN instead of the hardcoded 4096. The config echo reports it too.
| args.model_args = model_args | ||
|
|
||
|
|
||
| def _patch_trtllm_parse_logprobs(): |
There was a problem hiding this comment.
is this a bug with the official TRTLLM backend?
There was a problem hiding this comment.
Yes — a genuine bug in lm-eval 0.4.12's backend, confirmed both in source and at runtime.
TensorRT-LLM aligns prompt_logprobs to the next token. tensorrt_llm/executor/base_worker.py::_compute_pytorch_prompt_logprobs does:
# Pass prompt_token_ids with an offset of 1 for correct mapping to the context logits
prompt_token_ids = generation_result._generation_request.prompt_token_ids[1:] + first_generation_tokenso entry i is the distribution that predicted tokens[i + 1], and _topk_logprobs appends that token's id when it is not in the top-k. lm-eval instead reads prompt_logprobs[i][tokens[i]] and applies its own shift on top of that.
Probed directly against TRT-LLM 1.3.0rc23 with a 14-token prompt, for prompt_logprobs 0, 1 and 2 — tokens[i] is missing at every position, tokens[i+1] is present at every position. So it is not an edge case: it raises KeyError on the first request of every loglikelihood task (hellaswag, mmlu, arc, winogrande, ...). Only generate_until tasks work unpatched.
The corrected indexing now lives in lm_eval_trtllm.py and reproduces the removed implementation's hellaswag numbers exactly (0.7188 / 0.7812 at tp=1/2/4).
| _add_modelopt_args(run_parser) | ||
| args = cli.parse_args() | ||
| _inject_modelopt_args_into_model_args(args) | ||
| if getattr(args, "model", None) == "trtllm": |
There was a problem hiding this comment.
why we handle trtllm here? This should be hf only.
There was a problem hiding this comment.
Agreed — moved in dcedd37b4. lm_eval_hf.py is back to HF-only: the _patch_trtllm_parse_logprobs helper, its call site, and the logging import are all gone.
The override now lives in the new examples/llm_eval/lm_eval_trtllm.py, whose entire content is that one corrected _parse_logprobs plus cli_evaluate(). huggingface_example.sh and the docs call it instead.
| import copy | ||
| import gc | ||
| import logging | ||
| """[Deprecated] ModelOpt's TensorRT-LLM backend for lm-evaluation-harness. |
There was a problem hiding this comment.
I think it is ok to deprecate this file entirely
There was a problem hiding this comment.
Done — deleted outright in dcedd37b4, no translating shim.
examples/llm_eval/README.md and the CHANGELOG entry carry the replacement command:
python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=<tok>,checkpoint_dir=<ckpt>
-> python lm_eval_trtllm.py --model trtllm --model_args model=<ckpt>,tokenizer=<tok>,tensor_parallel_size=<tp>,max_input_len=4096.
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Design review (protocol applied). Problem: the repo maintained its own trt-llm lm-eval backend (lm_eval_tensorrt_llm.py, 213 lines on top of modelopt.deploy.llm.LLM); lm-eval 0.4.12 now ships lm_eval.models.trtllm_causallms, so the local copy is redundant. Deleting it in favour of upstream is the right direction — this PR removes a subsystem rather than adding one, and the net is -231/+152. Alternatives considered: (a) keep the local backend (rejected in the body with a feature-parity table — reasonable), (b) fix upstream and pin (the body says an upstream issue/PR is still needed and marks the shim as temporary — reasonable), (c) put the _parse_logprobs patch inside the existing lm_eval_hf.py entry point instead of adding a second entry script — the PR body actually describes option (c) as what was implemented, but the diff implements a separate lm_eval_trtllm.py. That gap is my main concern.
Findings:
-
The PR body no longer describes the diff, and the testing evidence is therefore unattributable. The body says
lm_eval_tensorrt_llm.py"becomes a deprecation shim", that "lm_eval_hf.pypatches_parse_logprobswhen--model trtllmis selected", that the usage ispython lm_eval_hf.py --model trtllm ..., that "the deprecation shim was run end to end", and answers "Is this change backward compatible?: ✅". The diff instead deletes the script (correctly reflected inCHANGELOG.rst, so the CHANGELOG and the body disagree) and adds a newlm_eval_trtllm.pythat callslm_eval.__main__.cli_evaluate. So the shipped code path — a flatpython lm_eval_trtllm.py --model trtllm --tasks ...invocation throughcli_evaluateon 0.4.12 — is not the one the accuracy tables were produced with. Please either update the body or re-state which command was actually run. -
cli_evaluatevs 0.4.12's subcommand CLI is unverified (CI-breaking if wrong).lm_eval_hf.pyhas to reach intoHarnessCLI._subparsers.choices["run"], i.e. 0.4.12's CLI is subcommand-based. It's not obvious thatlm_eval.__main__.cli_evaluatestill accepts flat--model/--tasks/--batch_sizeargs in 0.4.12; if it doesn't,test_qwen3_eval_fp8(the only coverage) fails. Related:huggingface_example.shforwards$lm_eval_flags, which can contain--trust_remote_code—lm_eval_hf.pygoes out of its way to move that flag intomodel_argsand null it out in the namespace, so please confirm the plain CLI accepts it and that it actually reaches the trtllm backend/tokenizer. -
No unit test for the off-by-one fix, which is the only new logic in the PR and is pure Python. A stub
outputsobject (SimpleNamespace(outputs=[SimpleNamespace(prompt_logprobs=[{tok: Logprob(logprob, rank)}...])])) would pin the alignment, theis_greedycomputation, and bothRuntimeErrorpaths without a GPU. Right now the fix is only exercised by a GPU/TRT-LLM example test whose reference numbers came from a different implementation (see 1). -
Minor: duplicated version guard between
lm_eval_hf.pyandlm_eval_trtllm.py, and the trtllm path loseslm_eval_hf.py's extras (--accuracy_lower_bound, trust-remote-code plumbing). Folding the patch intolm_eval_hf.pyunderargs.model == "trtllm"(as the body describes) would avoid both.
Licensing: new file carries the standard NVIDIA/Apache-2.0 header (year 2024 vs. LICENSE_HEADER's 2026 — cosmetic, matches existing files); no licensing surface. Doc/CHANGELOG references to the removed script are consistent — no dangling lm_eval_tensorrt_llm.py invocations remain in the repo.
|
|
||
| from importlib.metadata import version | ||
|
|
||
| from lm_eval.__main__ import cli_evaluate |
There was a problem hiding this comment.
Bot comment.
cli_evaluate on lm-eval 0.4.12: lm_eval_hf.py (same directory, same version floor) has to hook HarnessCLI._subparsers.choices["run"], which implies 0.4.12's CLI is subcommand-based. Has python lm_eval_trtllm.py --model trtllm --tasks ... --batch_size ... (flat args, no run subcommand) actually been run against 0.4.12? The PR body's usage/testing sections both reference lm_eval_hf.py --model trtllm plus a deprecation shim, neither of which exists in this diff, so the shipped entry point looks unvalidated — and test_qwen3_eval_fp8 is the only thing covering it.
There was a problem hiding this comment.
Yes — that exact flat-arg form is what produced the tp=1/2/4 numbers in the PR body. It works because 0.4.12's CLI keeps a legacy path in HarnessCLI.parse_args (lm_eval/_cli/harness.py:48-51):
if len(sys.argv) > 2 and sys.argv[1] not in self._subparsers.choices:
# Backward compatibility: arguments provided but no valid subcommand - insert 'run'
sys.argv.insert(1, "run")lm_eval_hf.py reaches for _subparsers.choices["run"] for a different reason — it adds the ModelOpt arguments (--quant_cfg, --calib_size, ...) to that subparser. lm_eval_trtllm.py adds no arguments, so plain cli_evaluate() is enough.
You were right that it looked unvalidated, though: the PR description was stale — it still documented lm_eval_hf.py --model trtllm and a deprecation shim, neither of which survived dcedd37b4. Rewritten now to match the shipped entry point, and 622b97c26 adds unit coverage so it is not resting on the GPU example test alone.
|
|
||
| return continuation_logprobs, is_greedy | ||
|
|
||
|
|
There was a problem hiding this comment.
Bot comment.
Two robustness points on the monkeypatch:
- It silently no-ops semantics if upstream changes: with only a
>=0.4.12floor and a<0.5pin, a 0.4.13 that fixes_parse_logprobswould be silently overridden by this patch (double-shifting again), and a rename/signature change would surface as an obscure failure. Consider asserting the attribute exists (and ideally a version ceiling / a comment tying removal to the upstream fix). - There is no test for
_parse_logprobsitself. This is the only new logic in the PR and needs no GPU — a stuboutputs(prompt_logprobslist of{token_id: Logprob(logprob, rank)}) can pin thei-1alignment, theis_greedy/rank != 1rule, themax(ctxlen, 1)edge case, and bothRuntimeErrorbranches. Please add one; the GPU example test's reference numbers were produced with a different implementation.
There was a problem hiding this comment.
Both taken, in 622b97c26.
Test — tests/examples/llm_eval/test_lm_eval_trtllm.py, no GPU and no tensorrt_llm install (the module import is guarded upstream, and _parse_logprobs is pure Python over the response object). Stubs prompt_logprobs as {token_id: Logprob(logprob, rank)} and covers the i-1 alignment, the rank != 1 -> is_greedy rule, the max(ctxlen, 1) edge, and both RuntimeError branches. Mutation-checked so the cases are load-bearing: dropping the -1 shift is caught by 5/5, ignoring ctxlen by 4/5.
Drift — agreed a silent double-shift on a fixed 0.4.13 is the dangerous case. I did not add a version ceiling, since that would also block unrelated 0.4.x fixes. Instead the module keeps the original in _UPSTREAM_PARSE_LOGPROBS and a test asserts it is still misaligned:
def test_upstream_is_still_misaligned():
"""Tripwire: when this fails, upstream fixed the bug and this file can be deleted."""
with pytest.raises(KeyError):
lm_eval_trtllm._UPSTREAM_PARSE_LOGPROBS(...)So the day upstream fixes it, CI fails and points at the removal rather than silently re-breaking the alignment. Plus a hasattr check that errors clearly if _parse_logprobs is renamed or removed.
One correction on the aside: the reference numbers were not produced by a different implementation. The tp=4 run was re-done against lm_eval_trtllm.py exactly as shipped after the code moved out of lm_eval_hf.py, and reproduces hellaswag 0.7188 / 0.7812. The unit test is still the right ask — it pins the alignment in a way an end-to-end accuracy number cannot.
| # 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" \ |
There was a problem hiding this comment.
Bot comment.
max_gen_toks=$BUILD_MAX_OUTPUT_LEN is carried over from the old backend, but the PR body states the upstream TRTLLM.__init__ only forwards a fixed set of kwargs and silently drops the rest. Is max_gen_toks in that set? If it is dropped, generation length is no longer capped by --output, which is exactly what test_qwen3_eval_fp8 relies on (output=128 "Cap generation length: gsm8k/humaneval otherwise generate up to 1024 tokens/sample") — the test would get slower and the updated comment would be misleading. Worth confirming, and dropping the arg if it's a no-op.
There was a problem hiding this comment.
Checked — max_gen_toks is honored, not dropped, so --output still caps generation and test_qwen3_eval_fp8's comment stays accurate.
It is an explicit named parameter of TRTLLM.__init__, not part of the **kwargs that get discarded (lm_eval/models/trtllm_causallms.py):
:60 max_gen_toks: int = 256, # named __init__ parameter
:86 self._max_gen_toks = max_gen_toks
:575 def max_gen_toks(self) -> int: return self._max_gen_toks
:411 kwargs, until, max_gen_toks = self.modify_gen_kwargs(
:412 gen_kwargs, eos=eos, default_max_gen_toks=self.max_gen_toks)
:427 SamplingParams(max_tokens=max_gen_toks, stop=until, **kwargs)
The PR body caused the confusion and I have fixed it: the "only forwards a fixed set" claim is about the kwargs handed to the TensorRT-LLM LLM API (tensor_parallel_size, max_input_len, kv_cache_config, ...), where engine-level extras really are silently dropped. lm-eval's own named parameters are consumed by the backend normally. Keeping the argument.
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of PR #2066 (head b68114bb). Design question is settled: this removes the 213-line local trt-llm backend in favour of lm-eval 0.4.12's built-in trtllm backend, and the maintainer explicitly asked for the override to live outside lm_eval_hf.py ("This should be hf only"), which the current diff does. Net size is small and confined to examples/ + tests/. Licensing: new files carry only the standard NVIDIA/Apache header — no licensing surface. No prompt-injection attempts in the PR content.
Previous critical findings — resolved
- 💬 Body/diff mismatch (shim vs. deleted script,
lm_eval_hf.py --model trtllmvs. new entry point): body rewritten to describe the shippedlm_eval_trtllm.py, CHANGELOG and README agree, no dangling references to the deleted script. - 💬 Flat-arg
cli_evaluateon 0.4.12's subcommand CLI: author citedlm_eval/_cli/harness.py:48-51(auto-insertsrunwhen argv[1] isn't a subcommand) and states the accuracy tables were produced with exactly that command form. Reasonable. - 💬 Missing unit test for the off-by-one:
tests/examples/llm_eval/test_lm_eval_trtllm.pyadds 6 cases pinning thei-1alignment,rank != 1→is_greedy, thectxlen=0edge, bothRuntimeErrorpaths, and an upstream-still-broken tripwire; follows the existingtest_cast_mxfp4_to_nvfp4.pysys.path convention. Silent-inflation fallback (continue) replaced by a raise plus a length guard, addressing the earlier eval-correctness concern. - Shell fixes landed:
tokenizer=$MODEL_ABS_PATH, quoted--model_args,LM_EVAL_TPoverride with the DeepEP/SM 12.0 rationale, and--input/BUILD_MAX_INPUT_LENwired intoparser.sh(previously echoed but never parsed — the fix is correct and defaulted to 4096, matching thetest_qwen3_eval_fp8comment update).
Still worth an owner look
- CI reach of the new tests.
example_tests.ymlrunstests/examples/llm_evalonly in the nightlytrtllm-non-prmatrix; the PR-gated trtllm job ishf_ptqonly. So these deliberately GPU-free tests — including the tripwire meant to catch a future 0.4.x that fixes upstream, and the assumption thatTRTLLM._parse_logprobsis a staticmethod with(tokens, outputs, ctxlen)— won't run on PRs. Placing them somewhere the PR gate collects would make the protection real. - 💬
--trust_remote_code— the author answered the flat-arg half of the earlier question but not this half.huggingface_example.shstill appends--trust_remote_codeto$lm_eval_flags, and it now flows into the plain lm-eval CLI, whereaslm_eval_hf.pygoes out of its way to rewrite it intomodel_argsand null it in the namespace. Behaviour matches the deleted script (also plaincli_evaluate), so it's not a regression, but nothing confirms it reaches the trtllm tokenizer, and CI doesn't exercise it. - Minor / optional: the
>=0.4.12version guard is now duplicated inlm_eval_hf.pyandlm_eval_trtllm.py; the trtllm path has no--accuracy_lower_boundequivalent; the guardlen(prompt_logprobs) < len(tokens) - 1still passes if exactly one prompt token were dropped (the author argues truncation can't occur on the loglikelihood path, which looks right); and the test only doesimportorskip("lm_eval")—importorskip("lm_eval.models.trtllm_causallms")would skip rather than error if upstream ever importstensorrt_llmeagerly.
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
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) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/examples/llm_eval/test_lm_eval_trtllm.py`:
- Around line 120-121: Move the imports in the affected test functions,
including datasets, EvaluatorConfig, and inspect, to module scope after the
existing optional lm_eval skip guard. If any optional or heavy dependency must
remain local, add a brief explanatory comment at each such import naming the
reason.
- Around line 112-132: Add a GPU-free test that instantiates the real TRTLLM
class, stubbing only AutoTokenizer.from_pretrained and tensorrt_llm.LLM while
capturing and asserting their forwarded trust_remote_code arguments. Extend
test_trust_remote_code_reaches_the_backend or add a focused neighboring test so
it exercises TRTLLM construction rather than stopping at
EvaluatorConfig._set_trust_remote_code. Move non-optional imports to module
scope and retain optional imports locally only with justification comments.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0248137d-df43-4006-9aed-58579d141b64
📒 Files selected for processing (1)
tests/examples/llm_eval/test_lm_eval_trtllm.py
| 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 | ||
|
|
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- test file outline ---'
ast-grep outline tests/examples/llm_eval/test_lm_eval_trtllm.py
printf '%s\n' '--- relevant test file ---'
cat -n tests/examples/llm_eval/test_lm_eval_trtllm.py | sed -n '1,190p'
printf '%s\n' '--- TRTLLM definitions and uses ---'
rg -n --glob '*.py' 'class TRTLLM|TRTLLM\(|trust_remote_code|AutoTokenizer|LLM\(' .
printf '%s\n' '--- contributing guidance ---'
rg -n -A8 -B4 'in-function imports|module.scope|mock|real implementation|examples' CONTRIBUTING.md tests/conftest.py 2>/dev/null || trueRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- changed files ---'
git diff --stat
printf '%s\n' '--- llm_eval example ---'
fd -t f . examples/llm_eval tests/examples/llm_eval | sort
printf '%s\n' '--- lm_eval package location ---'
python3 - <<'PY'
import importlib.util
spec = importlib.util.find_spec("lm_eval")
print(spec.origin if spec else "lm_eval not installed")
print(list(spec.submodule_search_locations or []) if spec else "")
PY
printf '%s\n' '--- installed TRTLLM implementation references ---'
python3 - <<'PY'
import importlib.util
spec = importlib.util.find_spec("lm_eval.models.trtllm_causallms")
print(spec.origin if spec else "module not installed")
PY
printf '%s\n' '--- focused source and dependency metadata ---'
rg -n --glob '*.py' --glob '*.toml' --glob '*.txt' --glob '*.md' \
'trtllm_causallms|TRTLLM|lm.eval|lm_eval|datasets' \
examples/llm_eval tests/examples/llm_eval pyproject.toml setup.py setup.cfg requirements* CONTRIBUTING.md 2>/dev/null | head -300Repository: NVIDIA/Model-Optimizer
Length of output: 1137
🌐 Web query:
NVIDIA Model-Optimizer lm_eval_trtllm.py test_lm_eval_trtllm.py TRTLLM trust_remote_code lm-evaluation-harness source
💡 Result:
In the context of evaluating models with TensorRT-LLM (TRTLLM), there are currently two primary, distinct approaches for integrating with evaluation frameworks: the official NVIDIA Model Optimizer integration and the native TRT-LLM backend within the EleutherAI lm-evaluation-harness. 1. NVIDIA Model Optimizer Integration: NVIDIA's Model Optimizer provides example scripts for evaluating quantized models [1][2]. The script lm_eval_tensorrt_llm.py acts as a wrapper to run TensorRT-LLM models within the evaluation workflow [3][1]. Users typically invoke it via command line: python lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=,checkpoint_dir= --tasks --batch_size <batch_size> [1][2]. 2. EleutherAI lm-evaluation-harness (Native Backend): As of version 0.4.12, the EleutherAI lm-evaluation-harness includes a native TRT-LLM backend (trt-llm) [4]. This was added to align feature parity with other backends like vLLM and sglang [5][6]. - Implementation: The backend is implemented in lm_eval/models/trtllm_causallms.py [6]. - Usage: You can use the native backend directly within the lm-evaluation-harness by specifying the trt-llm model type, which manages the TRT-LLM engine loading and execution internally [4][6]. Regarding your mention of test_lm_eval_trtllm.py, there is no widely documented or standard utility by that name in the primary NVIDIA Model Optimizer or EleutherAI repositories. It is likely a custom testing script for the evaluation wrapper. For general TensorRT-LLM evaluation, NVIDIA also provides the trtllm-eval tool, which is a unified entry point built on the TensorRT-LLM offline API and is distinct from the lm-evaluation-harness, though it may support similar tasks [7]. When using external tools like lm-evaluation-harness with TRT-LLM, users are sometimes advised to set arguments like add_special_tokens=False to prevent accuracy regressions caused by automatic BOS token insertion, depending on the specific model and task [8]. Trust_remote_code is a standard Hugging Face argument often required when loading custom model architectures, which can be passed through the model arguments in these frameworks if the underlying library supports it [6].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/llm_eval/README.md
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/llm_eval/README.md
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/llm_eval/lm_eval_tensorrt_llm.py
- 4: https://github.com/EleutherAI/lm-evaluation-harness/releases/tag/v0.4.12
- 5: feat: add TRT-LLM backend. EleutherAI/lm-evaluation-harness#3628
- 6: EleutherAI/lm-evaluation-harness@243c546
- 7: https://nvidia.github.io/TensorRT-LLM/commands/trtllm-eval.html
- 8: https://nvidia.github.io/TensorRT-LLM/latest/deployment-guide/deployment-guide-for-llama3.3-70b-on-trtllm.html
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from urllib.request import urlopen
urls = {
"trtllm_backend.py": "https://raw.githubusercontent.com/EleutherAI/lm-evaluation-harness/243c5463f67231b945f6add5f525e599c080ef63/lm_eval/models/trtllm_causallms.py",
"requirements.txt": "https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/main/examples/llm_eval/requirements.txt",
"contributing.md": "https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/main/CONTRIBUTING.md",
}
for name, url in urls.items():
print(f"--- {name} ---")
try:
text = urlopen(url, timeout=10).read().decode()
except Exception as exc:
print(f"FETCH_ERROR: {exc}")
continue
if name == "trtllm_backend.py":
lines = text.splitlines()
for i, line in enumerate(lines, 1):
if any(term in line for term in (
"class TRTLLM", "trust_remote_code", "AutoTokenizer",
"Tokenizer", "LLM(", "def __init__"
)):
start = max(1, i - 8)
end = min(len(lines), i + 12)
print(f"[{start}-{end}]")
print("\n".join(f"{j}: {lines[j-1]}" for j in range(start, end + 1)))
else:
print(text[:12000])
PYRepository: NVIDIA/Model-Optimizer
Length of output: 679
🏁 Script executed:
#!/bin/bash
set -e
for spec in \
"backend https://raw.githubusercontent.com/EleutherAI/lm-evaluation-harness/243c5463f67231b945f6add5f525e599c080ef63/lm_eval/models/trtllm_causallms.py" \
"requirements https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/main/examples/llm_eval/requirements.txt" \
"contributing https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/main/CONTRIBUTING.md"
do
name=${spec%% *}
url=${spec#* }
printf '%s\n' "--- $name ---"
curl -ksSfL "$url" | \
if [ "$name" = backend ]; then
nl -ba | grep -E -A12 -B8 'class TRTLLM|trust_remote_code|AutoTokenizer|Tokenizer|LLM\(|def __init__'
else
head -200
fi
doneRepository: NVIDIA/Model-Optimizer
Length of output: 313
🏁 Script executed:
#!/bin/bash
set -e
tmp=$(mktemp)
curl -kSfL \
'https://raw.githubusercontent.com/EleutherAI/lm-evaluation-harness/243c5463f67231b945f6add5f525e599c080ef63/lm_eval/models/trtllm_causallms.py' \
-o "$tmp"
grep -n -E -A12 -B8 'class TRTLLM|trust_remote_code|AutoTokenizer|Tokenizer|LLM\(|def __init__' "$tmp" || true
rm -f "$tmp"Repository: NVIDIA/Model-Optimizer
Length of output: 5043
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- llm_eval requirements ---'
cat -n examples/llm_eval/requirements.txt
printf '%s\n' '--- backend initialization source ---'
tmp=$(mktemp)
curl -kSfL \
'https://raw.githubusercontent.com/EleutherAI/lm-evaluation-harness/243c5463f67231b945f6add5f525e599c080ef63/lm_eval/models/trtllm_causallms.py' \
-o "$tmp"
sed -n '1,175p' "$tmp"
rm -f "$tmp"
printf '%s\n' '--- test guidance ---'
rg -n -A10 -B5 'imports at module|local import|optional|heavy|mock|real implementation|GPU-free|examples/llm_eval' \
CONTRIBUTING.md examples/llm_eval/README.md tests/examples/README.md tests/conftest.py 2>/dev/null || trueRepository: NVIDIA/Model-Optimizer
Length of output: 12537
Add a GPU-free TRTLLM forwarding test.
test_trust_remote_code_reaches_the_backend stops at EvaluatorConfig._set_trust_remote_code(), and test_trtllm_backend_accepts_trust_remote_code checks only the constructor signature. Instantiate the real TRTLLM, capture the arguments passed to AutoTokenizer.from_pretrained and tensorrt_llm.LLM, and stub only those external constructors. Move the non-optional test imports to module scope with justification comments for any remaining optional imports.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/examples/llm_eval/test_lm_eval_trtllm.py` around lines 112 - 132, Add a
GPU-free test that instantiates the real TRTLLM class, stubbing only
AutoTokenizer.from_pretrained and tensorrt_llm.LLM while capturing and asserting
their forwarded trust_remote_code arguments. Extend
test_trust_remote_code_reaches_the_backend or add a focused neighboring test so
it exercises TRTLLM construction rather than stopping at
EvaluatorConfig._set_trust_remote_code. Move non-optional imports to module
scope and retain optional imports locally only with justification comments.
Source: Path instructions
Remaining items from the re-reviewCI reach of the new tests — confirmed: Taken ( Answered, no change:
|
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of PR #2066 (examples/tests only, +331/−231). Design question stays settled: this removes the 213-line in-repo trt-llm lm-eval backend in favour of lm-eval 0.4.12's built-in trtllm backend, and the only new logic is a narrowly scoped _parse_logprobs override in a thin entry point — which is where the maintainer asked it to live ("this should be hf only"). No new subsystem, no licensing surface beyond the standard NVIDIA/Apache header on the new files, and no prompt-injection attempts in the PR content.
Previously-raised items now resolved
importorskip("lm_eval.models.trtllm_causallms")replaces the barelm_evalskip, so collection skips rather than errors if the backend ever importstensorrt_llmeagerly (d5f868c).- Off-by-one fix is unit-tested (
tests/examples/llm_eval/test_lm_eval_trtllm.py):i-1alignment,rank != 1→is_greedy,ctxlen=0edge, bothRuntimeErrorpaths, plus an upstream-still-broken tripwire; the silent-continuefallback is gone in favour of a hard raise. - Shell side is correct:
--input/BUILD_MAX_INPUT_LENnow actually parsed and defaulted inparser.sh(it was previously echoed only),tokenizer=$MODEL_ABS_PATH, quoted--model_args,LM_EVAL_TPescape hatch with the DeepEP/SM 12.0 rationale, and the lostfree_gpu_memory_fraction/expert-parallel tuning documented inexamples/llm_eval/README.md.
Still worth an owner look before sign-off
- 💬 CI reach — author decided to keep the new tests nightly-only (
llm_evalappears only in thetrtllm-non-prmatrix;trtllm-prishf_ptq), reasoning that the tripwire fires on new lm-eval releases rather than diffs and that addingllm_evaltotrtllm-prwould drag the 900s GPUtest_qwen3_eval_fp8onto every PR. Reasonable, but the consequence is that the deliberately GPU-free tests protecting the only new logic in this PR never gate a PR (ruff aside), and the tripwire also encodes a version-dependent assumption — that upstream_parse_logprobsis a staticmethod callable as(tokens=, outputs=, ctxlen=); if that changes it fails withTypeError, not the assertedKeyError, and only overnight. Worth a maintainer decision on whether these belong in a PR-run suite. - 💬
--trust_remote_code— author traced the full path (EvaluatorConfig._set_trust_remote_code→model_args["trust_remote_code"]→ namedTRTLLM.__init__param → tokenizer andllm_kwargs) and added three GPU-free tests. Flagging only because the coverage is static (config-level injection +inspect.signature), never exercisingTRTLLMconstruction, and no e2e run with atrust_remote_codemodel was done — acceptable to me, but it's a security-adjacent flag so an owner should accept the level of assurance. - Minor / convention:
tests/examples/llm_eval/test_lm_eval_trtllm.pystill has function-local imports (import datasets+from lm_eval.config.evaluate_config import EvaluatorConfigat ~120 and ~135,import inspectat ~148) with no stated justification; the module already has a top-levelimportorskipguard, so these can move to module scope after it. Not blocking. Also unaddressed-by-choice and reasonable: the duplicated>=0.4.12guard across the two entry points, no--accuracy_lower_boundon the trtllm path (parity with the deleted script), andlen(prompt_logprobs) < len(tokens) - 1tolerating a hypothetical one-token drop (the per-token raise covers it in practice).
| 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]) |
There was a problem hiding this comment.
[P1] Require the TensorRT-LLM version that provides this layout. This i - 1 mapping only became valid in TensorRT-LLM 1.3.0rc11. The hf_ptq README still recommends release:1.2.0; in that version _compute_pytorch_prompt_logprobs does not pass prompt token IDs into compute_logprobs, so prompt_logprobs=1 retains only top-1. A normal non-greedy continuation token is therefore absent and this branch raises, breaking likelihood tasks that worked through the deleted wrapper. See v1.2.0 versus v1.3.0rc11. Please add an explicit >=1.3.0rc11 runtime guard and update the documented container requirement, or keep a compatible path for the supported 1.1/1.2 releases.
There was a problem hiding this comment.
Confirmed and fixed in 6a54747c3. I diffed the two tags rather than taking it on trust, and you are right:
# v1.2.0 -- no prompt token ids
logprobs_result = compute_logprobs(logprob_params.prompt_logprobs, None,
context_logits, None, None)
# v1.3.0rc11 -- prompt token ids passed in
prompt_token_ids = generation_result._generation_request.prompt_token_ids[1:] + first_generation_token
logprobs_result = compute_logprobs(logprob_params.prompt_logprobs, None,
context_logits, None, None,
prompt_token_ids)With tokens=None, _topk_logprobs never reaches its "append the requested token if it is not in top-k" branch, so prompt_logprobs=1 keeps only the argmax. A non-greedy continuation token is genuinely absent -- there is no correct value to recover, for this override or for lm-eval's own _parse_logprobs. And since hf_ptq/README.md recommends release:1.2.0, a user following the docs would hit exactly this.
Guard placement. I put it in _parse_logprobs, not at startup. A blanket startup abort would break generative tasks (gsm8k, ifeval) on 1.1/1.2, which work fine on the old layout and worked through the deleted wrapper — that would be a fresh regression rather than a fix. So:
_check_trtllm_version()raises from_parse_logprobson the first loglikelihood request, checked once.__main__calls the same function and prints a warning, so an unusable container is visible immediately rather than after a multi-minute model load.
Net: loglikelihood on <1.3.0rc11 fails with an explicit "use a newer container" message instead of a confusing missing-token error; generative tasks are untouched.
Docs, as you asked: examples/llm_eval/README.md and examples/hf_ptq/README.md (right under the release:1.2.0 recommendation) both state the requirement, and the CHANGELOG entry does too.
Tests: seven cases — reject on 1.1.0rc2 / 1.2.0 / 1.3.0rc10, accept on 1.3.0rc11 / 1.3.0rc23 / 1.3.0 / 1.4.0, plus one asserting the guard fires from _parse_logprobs and not only from __main__. A missing tensorrt_llm is treated as "nothing to check", which keeps the suite install-free; that branch is unreachable in a real run because the backend refuses to build a model without it.
| import datasets | ||
| from lm_eval.config.evaluate_config import EvaluatorConfig | ||
|
|
||
| monkeypatch.setattr(datasets.config, "HF_DATASETS_TRUST_REMOTE_CODE", False) |
There was a problem hiding this comment.
[P1] These tests fail with supported datasets 4.x because HF_DATASETS_TRUST_REMOTE_CODE is not defined until lm-eval assigns it. monkeypatch.setattr(..., raising=True) therefore raises AttributeError before testing the production behavior. I reproduced this with exact lm_eval==0.4.12 and datasets==4.8.4: 7 tests passed and these two trust-remote-code tests failed. Please use raising=False or defensively establish/remove the attribute. The relevant trtllm-non-pr suite was skipped by the current PR run, which is why the checks remain green.
There was a problem hiding this comment.
Reproduced and fixed in 6a54747c3. I removed HF_DATASETS_TRUST_REMOTE_CODE from datasets.config to simulate 4.x and got exactly your result — 7 passed, 2 failed with AttributeError from monkeypatch.setattr, before either test exercised anything.
My venv has datasets 3.3.2, which still defines the attribute. That is the whole reason it was green locally and red for you; thanks for pinning the version, it made this a two-minute repro instead of a guess.
Fixed with raising=False on both calls, which is the right form here: lm-eval creates the attribute on assignment, and monkeypatch still removes it again on teardown, so the isolation the tests wanted is preserved either way.
Now verified under both: 17 passed on datasets 3.3.2, and 17 passed with the attribute deleted. Also re-ran through the shared tests/examples/conftest.py rather than --noconftest, so the real collection path is covered.
And your closing point is well taken — the trtllm-non-pr suite being skipped is exactly why this reached you instead of CI. That is the nightly-only gap discussed above; both of these findings would have been caught by it.
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
| 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 |
There was a problem hiding this comment.
Maybe we don't need this change
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
3 participants
What does this PR do?
Type of change: documentation / example update (with a behaviour fix)
lm-evaluation-harness 0.4.12 is the first release that ships a TensorRT-LLM backend
(
lm_eval.models.trtllm_causallms, registered astrtllm) — it is absent in 0.4.10 and0.4.11. This example no longer maintains its own, so:
lm_eval[api,ifeval]>=0.4.12,<0.5(the 0.5.0.dev line drops the file) and bumplm_eval_hf.py's version guard to match.examples/llm_eval/lm_eval_tensorrt_llm.py(thetrt-llmmodel). Replacepython lm_eval_tensorrt_llm.py --model trt-llm --model_args tokenizer=<tok>,checkpoint_dir=<ckpt>with
python lm_eval_trtllm.py --model trtllm --model_args model=<ckpt>,tokenizer=<tok>.examples/llm_eval/lm_eval_trtllm.py, whose entire content is one corrected_parse_logprobspluscli_evaluate()(see below).lm_eval_hf.pystays HF-only.examples/hf_ptq/scripts/huggingface_example.shand the docs use the upstream backend.parser.shgains--input(BUILD_MAX_INPUT_LEN, default 4096) — it already echoedthat variable but never parsed or defaulted it, so it printed empty on every run.
Why
lm_eval_trtllm.pyexists: an upstream off-by-oneTensorRT-LLM aligns
prompt_logprobsto the next token.executor/base_worker.py:So entry
iis the distribution that predictedtokens[i + 1], and_topk_logprobsappends that token's id when it is not in the top-k. lm-eval's
_parse_logprobsinsteadreads
prompt_logprobs[i][tokens[i]]and applies its own shift on top, which raisesKeyErroron the first request of every loglikelihood task (hellaswag, mmlu, arc, ...):Probed against TRT-LLM 1.3.0rc23 with a 14-token prompt for
prompt_logprobs0, 1 and 2:tokens[i]is missing at every position,tokens[i+1]is present at every position.Only
generate_untiltasks work unpatched. This wants an upstream issue againstEleutherAI/lm-evaluation-harness.
The override also fails loudly rather than quietly: it checks
prompt_logprobscoversevery prompt token and raises on a missing token, instead of skipping the term and
silently inflating the reported accuracy.
Defaults that must be set explicitly
TRTLLM.__init__accepts**kwargsbut forwards only a fixed set to the TensorRT-LLMLLMAPI, so extra--model_argsaimed at the engine are silently dropped. (lm-eval'sown named parameters —
max_gen_toks,batch_size,truncation_side, ... — are honorednormally.) Two engine defaults are unsafe for few-shot eval:
tensor_parallel_sizedefaults to 1 (the deleted wrapper used every visible GPU).max_input_lendefaults to 2048, and longer prompts are silently left-truncated —5-shot MMLU/gsm8k prompts exceed that.
Usage
python lm_eval_trtllm.py --model trtllm \ --model_args model=<quantized checkpoint dir>,tokenizer=<HF model folder>,tensor_parallel_size=<tp>,max_batch_size=<bs>,max_input_len=4096,max_output_len=512 \ --tasks hellaswag,gsm8k \ --batch_size <bs>Flat arguments (no
runsubcommand) are what 0.4.12'sHarnessCLI.parse_argsinsertsrunfor automatically (_cli/harness.py:48-51); this is the exact command form used forthe results below.
Testing
Unit —
tests/examples/llm_eval/test_lm_eval_trtllm.py, no GPU and notensorrt_llminstall: stubs the response object and pins the
i-1alignment, therank != 1→is_greedyrule, thectxlen=0edge, and bothRuntimeErrorpaths. Mutation-checked —dropping the
-1shift is caught by 5/5 cases, ignoringctxlenby 4/5. A sixth test is atripwire: it asserts lm-eval's own implementation is still misaligned, so a future
0.4.x that fixes the bug fails the test and says to delete this file rather than being
silently re-broken by the override.
End to end —
nvidia/Qwen3.5-122B-A10B-NVFP4(NVFP4 MoE, 256 experts) on 4x B300,TRT-LLM 1.3.0rc23, lm-eval 0.4.12,
--limit 32:trt-llm), tp=4lm_eval_trtllm.py, tp=1lm_eval_trtllm.py, tp=2lm_eval_trtllm.py, tp=4to the deleted implementation — the alignment fix is exact, not approximate.
stop=sequences and per-request
SamplingParams; the old wrapper used beam-search-of-1 withpost-hoc string truncation).
KeyError.lm_eval_hf.pyintolm_eval_trtllm.py.Note: NVFP4 fused-MoE has no CUTLASS tactic on Hopper (
No supported MoE GEMM tactic remains after replacing unsupported NO_SMEM epilogues.), so this had to be validated onBlackwell.
Feature parity notes
Gained from upstream:
loglikelihood_rolling(wasNotImplementedError), pipelineparallelism,
add_bos_tokenauto-detection, prompt truncation, per-request sampling params,prompt_logprobsinstead of full-vocab context logits (much lower memory), thinking-taghandling,
batch_size=auto.Not reachable through the upstream backend (were set by
modelopt.deploy.llm.LLM):enable_attention_dpfor MoE,CudaGraphConfig,enable_chunked_prefill,moe_expert_parallel_size=1, andfree_gpu_memory_fraction=0.7with a cappedkv_cache.max_tokens— upstream uses the TRT-LLM default 0.9 (observed allocating 218 GiBof paged KV cache on B300), so OOM risk is higher on smaller GPUs. This is documented in
examples/llm_eval/README.md, andhuggingface_example.shhonours a presetLM_EVAL_TPso users can lower the tensor-parallel size without editing the script.
Before your PR is "Ready for review"
lm_eval_tensorrt_llm.pyis removed and the CLI changes (--model trt-llm→trtllm,checkpoint_dir=→model=). Migration command is in the README and CHANGELOG.CONTRIBUTING.md: ✅ — no new dependency; existinglm_evalpin tightened.tests/examples/llm_eval/test_lm_eval_trtllm.py(6 cases, no GPU).dcedd37b4and622b97c26.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
trtllmbackend.Documentation
Deprecations
Updates