From ef038b0caeedafbd3076b5e43659d5a43e7f9c39 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 22 Jul 2026 11:50:03 -0700 Subject: [PATCH 1/8] [None][feat] Log running metric estimates during long lm-eval runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit lm-eval computes metrics only after every response is collected, so a multi-hour trtllm-eval run gives no quality signal until the very end — there is no way to tell whether a long job is on track or should be killed early. Add _RunningScoreTracker to LmEvalWrapper.generate_until: as each response completes, apply the owning task's filter ensembles and process_results to a throwaway copy of the instance and log a running mean every N responses, e.g. Partial scores after 200/1319 responses (estimate, 0~100): exact_match,strict-match ~ 93.50 | exact_match,flexible-extract ~ 94.00 Opt-in via TLLM_EVAL_PARTIAL_SCORES_EVERY=N (unset/0 = off, no behavior change). Env-var driven so it uniformly covers every lm-eval-backed task without per-task CLI plumbing. Scoring happens on copies — harness state and the final score are unaffected — and any scoring failure permanently disables the tracker for the run instead of failing the eval. Following lm-eval's calling convention, the filtered response is passed to process_results wrapped in a list (one entry per repeat); a regression test pins this down since a bare string would silently score only the first character. Estimates are meaningful for per-sample-decomposable metrics (e.g. exact_match); the multimodal wrapper's generate_until override does not implement partial scoring yet. Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 134 +++++++++++++++++++++++++- tests/unittest/others/test_lm_eval.py | 132 +++++++++++++++++++++++++ 2 files changed, 261 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index b077747b6c03..2150d7d79d8a 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -15,6 +15,7 @@ import copy import json import os +from collections import defaultdict from contextlib import contextmanager from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple @@ -50,6 +51,87 @@ # https://github.com/EleutherAI/lm-evaluation-harness/blob/7f04db12d2f8e7a99a0830d99eb78130e1ba2122/lm_eval/models/hf_vlms.py#L25 LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER = "" +# Interval (in completed responses) for logging running metric estimates +# during long evals; see _RunningScoreTracker. 0/unset disables the feature. +PARTIAL_SCORES_ENV_VAR = "TLLM_EVAL_PARTIAL_SCORES_EVERY" + + +class _RunningScoreTracker: + """Best-effort running metric estimates over completed eval responses. + + lm-eval computes metrics only after every response is collected, so a + multi-hour eval gives no quality signal until the very end. This tracker + scores each response as it completes — applying the owning task's filter + ensembles and ``process_results`` to a throwaway copy of the instance — + and logs a running aggregate every ``interval`` responses, so an operator + can tell whether a long job is on track or should be killed early. + + The running numbers are estimates: per-sample filtering happens outside + the harness's batch path, and only per-sample-decomposable metrics (e.g. + exact_match) aggregate meaningfully as a mean. The final score reported + by lm-eval is unaffected — scoring here happens on copies. Any failure + (exotic task/filter/metric shapes) permanently disables the tracker for + the run and never fails the eval itself. + """ + + def __init__(self, task_dict: dict, interval: int): + self.interval = interval + self.tasks = {} + self._collect_tasks(task_dict) + self.metric_sums = defaultdict(float) + self.metric_counts = defaultdict(int) + self.disabled = False + + def _collect_tasks(self, task_dict: dict) -> None: + for name, obj in task_dict.items(): + if isinstance(obj, dict): # task group + self._collect_tasks(obj) + else: + self.tasks[name] = obj + + def update(self, instance, text: str) -> None: + if self.disabled: + return + try: + task = self.tasks.get(instance.task_name) + if task is None or not getattr(task, "_filters", None): + raise ValueError( + f"no scorable task for {instance.task_name!r}") + # Score a shallow copy: the harness appends resps / applies + # filters to the real instance later, and must see it untouched. + probe = copy.copy(instance) + probe.resps = [text] + probe.filtered_resps = {} + for ensemble in task._filters: + ensemble.apply([probe]) + # lm-eval evaluator passes a list of filtered responses (one + # per repeat/request), so process_results does results[0] to + # get the prediction. Match that contract here — wrapping in + # a list gives the same calling convention. + metrics = task.process_results( + probe.doc, [probe.filtered_resps[ensemble.name]]) + for metric, value in metrics.items(): + if isinstance(value, (int, float)): + key = f"{metric},{ensemble.name}" + self.metric_sums[key] += value + self.metric_counts[key] += 1 + except Exception as e: + logger.info( + f"Partial scoring disabled for this run ({type(e).__name__}: {e})" + ) + self.disabled = True + + def maybe_log(self, done: int, total: int) -> None: + if self.disabled or not self.metric_counts: + return + if done % self.interval != 0 and done != total: + return + stats = " | ".join( + f"{key} ~ {100 * self.metric_sums[key] / count:.2f}" + for key, count in self.metric_counts.items()) + logger.info(f"Partial scores after {done}/{total} responses " + f"(estimate, 0~100): {stats}") + class LmEvalWrapper(TemplateLM): @@ -62,7 +144,9 @@ def __init__(self, is_force_single_image: bool = False, output_dir: Optional[str] = None, sampling_override: bool = False, - preserve_caller_max_tokens: bool = False): + preserve_caller_max_tokens: bool = False, + partial_scores_every: Optional[int] = None, + partial_scoring_task_dict: Optional[dict] = None): super().__init__() self.llm = llm self.sampling_params = sampling_params @@ -77,6 +161,11 @@ def __init__(self, # task yaml's max_gen_toks. Opt-in for thinking models (e.g. Kimi K2.5) # whose chain-of-thought output exceeds lm-eval's default (~512). self.preserve_caller_max_tokens = preserve_caller_max_tokens + # When set (with the task_dict), log running metric estimates every + # N completed responses during generate_until — a liveness/quality + # signal for long evals. See _RunningScoreTracker. + self.partial_scores_every = partial_scores_every + self.partial_scoring_task_dict = partial_scoring_task_dict @property def eot_token_id(self) -> int: @@ -171,11 +260,19 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: streaming=self.streaming) results.append(output) + scorer = None + if self.partial_scores_every and self.partial_scoring_task_dict: + scorer = _RunningScoreTracker(self.partial_scoring_task_dict, + self.partial_scores_every) + outputs = [] - for output in tqdm(results, - desc="Fetching responses", - disable=disable_tqdm): + for output, request in zip( + tqdm(results, desc="Fetching responses", + disable=disable_tqdm), requests): outputs.append(output.result()) + if scorer is not None: + scorer.update(request, outputs[-1].outputs[0].text) + scorer.maybe_log(len(outputs), len(requests)) if self.output_dir: dump_inference_results(self.output_dir, outputs, @@ -208,7 +305,9 @@ def __init__(self, output_dir: Optional[str] = None, sampling_override: bool = False, preserve_caller_max_tokens: bool = False, - post_process_fn: Optional[Callable[[str], str]] = None): + post_process_fn: Optional[Callable[[str], str]] = None, + partial_scores_every: Optional[int] = None, + partial_scoring_task_dict: Optional[dict] = None): """ Initialize the multimodal wrapper. @@ -227,6 +326,10 @@ def __init__(self, to model outputs before scoring. Used by Kimi K2.5 to strip ``...`` and extract the final answer (see ``tensorrt_llm.evaluate.post_processing``). + partial_scores_every: Accepted for interface parity with + LmEvalWrapper; the multimodal generate_until override does + not implement partial scoring yet. + partial_scoring_task_dict: See partial_scores_every. """ super().__init__( llm, @@ -238,6 +341,8 @@ def __init__(self, output_dir=output_dir, sampling_override=sampling_override, preserve_caller_max_tokens=preserve_caller_max_tokens, + partial_scores_every=partial_scores_every, + partial_scoring_task_dict=partial_scoring_task_dict, ) # NOTE: Required by lm_eval to identify this as a multimodal model @@ -611,6 +716,22 @@ def evaluate(self, import lm_eval lm_cls = MultimodalLmEvalWrapper if self.MULTIMODAL else LmEvalWrapper + # Opt-in running metric estimates for long evals: log a partial + # score every N completed responses (see _RunningScoreTracker). + # Env-var driven so it uniformly covers every lm-eval-backed task + # without per-task CLI plumbing. + partial_scores_every = None + env_interval = os.environ.get(PARTIAL_SCORES_ENV_VAR) + if env_interval: + try: + partial_scores_every = int(env_interval) + except ValueError: + raise ValueError( + f"{PARTIAL_SCORES_ENV_VAR} must be an integer, got " + f"{env_interval!r}") from None + if partial_scores_every <= 0: + partial_scores_every = None + lm_kwargs: Dict[str, Any] = dict( sampling_params=sampling_params, streaming=streaming, @@ -619,6 +740,9 @@ def evaluate(self, is_force_single_image=is_force_single_image, output_dir=self.output_dir, sampling_override=sampling_override, + partial_scores_every=partial_scores_every, + partial_scoring_task_dict=self.task_dict + if partial_scores_every else None, ) # post_process_fn / preserve_caller_max_tokens only consumed by multimodal. if self.MULTIMODAL: diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index b14128717be7..a3a0be55dac5 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -721,3 +721,135 @@ def test_mode_cot_included_in_example_format(): finally: # Restore module state for other tests running in the same session. _reload_mmmu_pro_utils(None) + + +# =========================================================================== +# _RunningScoreTracker — partial score estimates during generate_until +# =========================================================================== +# +# Enabled via TLLM_EVAL_PARTIAL_SCORES_EVERY; scores each completed response +# with the owning task's filters + process_results on a throwaway instance +# copy, and must never disturb the real instance or fail the eval. + + +class _FakeEnsemble: + """Minimal stand-in for lm_eval.api.filter.FilterEnsemble.""" + + def __init__(self, name): + self.name = name + + def apply(self, instances): + for inst in instances: + # Trivial "take_first" pipeline. + inst.filtered_resps[self.name] = inst.resps[0] + + +class _FakeTask: + + def __init__(self): + self._filters = [_FakeEnsemble("strict-match")] + + def process_results(self, doc, results): + # Mirror lm-eval's ConfigurableTask.process_results: results is a list + # (one entry per repeat/request), and the prediction is results[0]. + return {"exact_match": float(results[0] == doc["answer"])} + + +class _FakeInstance: + + def __init__(self, task_name, doc): + self.task_name = task_name + self.doc = doc + self.resps = [] + self.filtered_resps = {} + + +def _make_tracker(interval=2): + from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + return _RunningScoreTracker({"fake_task": _FakeTask()}, interval) + + +def test_running_score_tracker_aggregates_mean(): + """Running estimate is the mean of per-sample metric values.""" + tracker = _make_tracker() + docs = [{"answer": "42"}, {"answer": "7"}, {"answer": "1"}] + responses = ["42", "0", "1"] # right, wrong, right + for doc, text in zip(docs, responses): + tracker.update(_FakeInstance("fake_task", doc), text) + assert not tracker.disabled + key = "exact_match,strict-match" + assert tracker.metric_counts[key] == 3 + assert tracker.metric_sums[key] == 2.0 + + +def test_running_score_tracker_does_not_mutate_instance(): + """The real instance stays untouched — the harness fills it in later.""" + tracker = _make_tracker() + instance = _FakeInstance("fake_task", {"answer": "42"}) + tracker.update(instance, "42") + assert instance.resps == [] + assert instance.filtered_resps == {} + + +def test_running_score_tracker_unknown_task_disables(): + """Any scoring failure permanently disables the tracker, never raises.""" + tracker = _make_tracker() + tracker.update(_FakeInstance("unknown_task", {"answer": "42"}), "42") + assert tracker.disabled + # Subsequent updates and logging are silent no-ops. + tracker.update(_FakeInstance("fake_task", {"answer": "42"}), "42") + assert not tracker.metric_counts + tracker.maybe_log(10, 100) + + +def test_running_score_tracker_logs_on_interval(caplog): + """maybe_log emits at every `interval` responses and at completion.""" + tracker = _make_tracker(interval=2) + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + tracker.update(_FakeInstance("fake_task", {"answer": "1"}), "1") + tracker.maybe_log(1, 3) # off-interval, not final -> no log + mock_logger.info.assert_not_called() + tracker.update(_FakeInstance("fake_task", {"answer": "1"}), "0") + tracker.maybe_log(2, 3) # on-interval -> logs + assert mock_logger.info.call_count == 1 + tracker.update(_FakeInstance("fake_task", {"answer": "1"}), "1") + tracker.maybe_log(3, 3) # final response -> logs + assert mock_logger.info.call_count == 2 + message = mock_logger.info.call_args[0][0] + assert "2/3" not in message # latest call reports 3/3 + assert "3/3" in message + assert "exact_match,strict-match" in message + # 2 of 3 correct -> ~66.67 on the 0~100 scale. + assert "66.67" in message + + +def test_running_score_tracker_process_results_list_convention(): + """process_results receives a list, not a bare string (regression for GSM8K bug). + + lm-eval's ConfigurableTask.process_results does ``result = results[0]`` to + extract the prediction from the list of per-repeat responses. If the tracker + passes the filtered_resp string directly instead of wrapping it in a list, + ``results[0]`` silently returns the *first character* of the string, causing + multi-digit answers to score as misses (~26% on GSM8K) while single-digit + answers accidentally match. + """ + tracker = _make_tracker() + # Use a multi-digit answer so the first-character bug is observable: + # "42" would produce results[0]=="4" if the list wrap were missing. + doc = {"answer": "42"} + tracker.update(_FakeInstance("fake_task", doc), "42") + assert not tracker.disabled + key = "exact_match,strict-match" + assert tracker.metric_sums[key] == 1.0, ( + "multi-digit answer scored as miss — process_results likely received " + "a bare string so results[0] returned only the first character" + ) + + +def test_running_score_tracker_task_groups_flattened(): + """Nested task_dict groups resolve to their leaf tasks.""" + from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + tracker = _RunningScoreTracker({"group": {"fake_task": _FakeTask()}}, 1) + tracker.update(_FakeInstance("fake_task", {"answer": "42"}), "42") + assert not tracker.disabled + assert tracker.metric_counts["exact_match,strict-match"] == 1 From 31c468e0cc2c3cb5df2d9fedf2253104a737a184 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 22 Jul 2026 14:22:33 -0500 Subject: [PATCH 2/8] [None][feat] Address CodeRabbit review on partial lm-eval scoring - Include task_name in metric aggregation key so multi-task evals keep per-task running estimates separate - Add type annotations to _RunningScoreTracker (__init__ -> None, instance: Any in update) - Extract env-var parsing into _parse_partial_scores_env() to make it unit-testable - Add tests: task-key separation, env-var parsing (positive/zero/negative/ invalid/unset), and generate_until tracker invocation Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 43 +++++----- tests/unittest/others/test_lm_eval.py | 111 ++++++++++++++++++++++++-- 2 files changed, 126 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 2150d7d79d8a..16003e77385d 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -74,7 +74,7 @@ class _RunningScoreTracker: the run and never fails the eval itself. """ - def __init__(self, task_dict: dict, interval: int): + def __init__(self, task_dict: dict, interval: int) -> None: self.interval = interval self.tasks = {} self._collect_tasks(task_dict) @@ -89,14 +89,13 @@ def _collect_tasks(self, task_dict: dict) -> None: else: self.tasks[name] = obj - def update(self, instance, text: str) -> None: + def update(self, instance: Any, text: str) -> None: if self.disabled: return try: task = self.tasks.get(instance.task_name) if task is None or not getattr(task, "_filters", None): - raise ValueError( - f"no scorable task for {instance.task_name!r}") + raise ValueError(f"no scorable task for {instance.task_name!r}") # Score a shallow copy: the harness appends resps / applies # filters to the real instance later, and must see it untouched. probe = copy.copy(instance) @@ -112,7 +111,7 @@ def update(self, instance, text: str) -> None: probe.doc, [probe.filtered_resps[ensemble.name]]) for metric, value in metrics.items(): if isinstance(value, (int, float)): - key = f"{metric},{ensemble.name}" + key = f"{instance.task_name},{metric},{ensemble.name}" self.metric_sums[key] += value self.metric_counts[key] += 1 except Exception as e: @@ -126,13 +125,25 @@ def maybe_log(self, done: int, total: int) -> None: return if done % self.interval != 0 and done != total: return - stats = " | ".join( - f"{key} ~ {100 * self.metric_sums[key] / count:.2f}" - for key, count in self.metric_counts.items()) + stats = " | ".join(f"{key} ~ {100 * self.metric_sums[key] / count:.2f}" + for key, count in self.metric_counts.items()) logger.info(f"Partial scores after {done}/{total} responses " f"(estimate, 0~100): {stats}") +def _parse_partial_scores_env() -> Optional[int]: + """Parse TLLM_EVAL_PARTIAL_SCORES_EVERY and return the logging interval or None.""" + env_interval = os.environ.get(PARTIAL_SCORES_ENV_VAR) + if not env_interval: + return None + try: + value = int(env_interval) + except ValueError: + raise ValueError(f"{PARTIAL_SCORES_ENV_VAR} must be an integer, got " + f"{env_interval!r}") from None + return value if value > 0 else None + + class LmEvalWrapper(TemplateLM): def __init__(self, @@ -267,8 +278,8 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: outputs = [] for output, request in zip( - tqdm(results, desc="Fetching responses", - disable=disable_tqdm), requests): + tqdm(results, desc="Fetching responses", disable=disable_tqdm), + requests): outputs.append(output.result()) if scorer is not None: scorer.update(request, outputs[-1].outputs[0].text) @@ -720,17 +731,7 @@ def evaluate(self, # score every N completed responses (see _RunningScoreTracker). # Env-var driven so it uniformly covers every lm-eval-backed task # without per-task CLI plumbing. - partial_scores_every = None - env_interval = os.environ.get(PARTIAL_SCORES_ENV_VAR) - if env_interval: - try: - partial_scores_every = int(env_interval) - except ValueError: - raise ValueError( - f"{PARTIAL_SCORES_ENV_VAR} must be an integer, got " - f"{env_interval!r}") from None - if partial_scores_every <= 0: - partial_scores_every = None + partial_scores_every = _parse_partial_scores_env() lm_kwargs: Dict[str, Any] = dict( sampling_params=sampling_params, diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index a3a0be55dac5..858df18f7e63 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -29,6 +29,8 @@ import os from unittest.mock import MagicMock, patch +import pytest + from tensorrt_llm.evaluate.covost2 import CoVoST2 from tensorrt_llm.evaluate.lm_eval import ( LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER, @@ -745,7 +747,6 @@ def apply(self, instances): class _FakeTask: - def __init__(self): self._filters = [_FakeEnsemble("strict-match")] @@ -756,7 +757,6 @@ def process_results(self, doc, results): class _FakeInstance: - def __init__(self, task_name, doc): self.task_name = task_name self.doc = doc @@ -766,6 +766,7 @@ def __init__(self, task_name, doc): def _make_tracker(interval=2): from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + return _RunningScoreTracker({"fake_task": _FakeTask()}, interval) @@ -777,7 +778,7 @@ def test_running_score_tracker_aggregates_mean(): for doc, text in zip(docs, responses): tracker.update(_FakeInstance("fake_task", doc), text) assert not tracker.disabled - key = "exact_match,strict-match" + key = "fake_task,exact_match,strict-match" assert tracker.metric_counts[key] == 3 assert tracker.metric_sums[key] == 2.0 @@ -802,7 +803,7 @@ def test_running_score_tracker_unknown_task_disables(): tracker.maybe_log(10, 100) -def test_running_score_tracker_logs_on_interval(caplog): +def test_running_score_tracker_logs_on_interval(): """maybe_log emits at every `interval` responses and at completion.""" tracker = _make_tracker(interval=2) with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: @@ -818,7 +819,7 @@ def test_running_score_tracker_logs_on_interval(caplog): message = mock_logger.info.call_args[0][0] assert "2/3" not in message # latest call reports 3/3 assert "3/3" in message - assert "exact_match,strict-match" in message + assert "fake_task,exact_match,strict-match" in message # 2 of 3 correct -> ~66.67 on the 0~100 scale. assert "66.67" in message @@ -839,7 +840,7 @@ def test_running_score_tracker_process_results_list_convention(): doc = {"answer": "42"} tracker.update(_FakeInstance("fake_task", doc), "42") assert not tracker.disabled - key = "exact_match,strict-match" + key = "fake_task,exact_match,strict-match" assert tracker.metric_sums[key] == 1.0, ( "multi-digit answer scored as miss — process_results likely received " "a bare string so results[0] returned only the first character" @@ -849,7 +850,103 @@ def test_running_score_tracker_process_results_list_convention(): def test_running_score_tracker_task_groups_flattened(): """Nested task_dict groups resolve to their leaf tasks.""" from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + tracker = _RunningScoreTracker({"group": {"fake_task": _FakeTask()}}, 1) tracker.update(_FakeInstance("fake_task", {"answer": "42"}), "42") assert not tracker.disabled - assert tracker.metric_counts["exact_match,strict-match"] == 1 + assert tracker.metric_counts["fake_task,exact_match,strict-match"] == 1 + + +def test_running_score_tracker_separate_keys_per_task(): + """Two tasks with the same metric/filter don't mix their running estimates.""" + from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + + task_a = _FakeTask() + task_b = _FakeTask() + tracker = _RunningScoreTracker({"task_a": task_a, "task_b": task_b}, 999) + tracker.update(_FakeInstance("task_a", {"answer": "x"}), "x") # correct + tracker.update(_FakeInstance("task_b", {"answer": "x"}), "y") # wrong + assert not tracker.disabled + assert tracker.metric_sums["task_a,exact_match,strict-match"] == 1.0 + assert tracker.metric_sums["task_b,exact_match,strict-match"] == 0.0 + + +# =========================================================================== +# _parse_partial_scores_env — env-var parsing +# =========================================================================== + + +def test_parse_partial_scores_env_positive(monkeypatch): + """A positive integer returns that interval.""" + from tensorrt_llm.evaluate.lm_eval import PARTIAL_SCORES_ENV_VAR, _parse_partial_scores_env + + monkeypatch.setenv(PARTIAL_SCORES_ENV_VAR, "100") + assert _parse_partial_scores_env() == 100 + + +def test_parse_partial_scores_env_zero_disables(monkeypatch): + """Zero disables partial scoring (returns None).""" + from tensorrt_llm.evaluate.lm_eval import PARTIAL_SCORES_ENV_VAR, _parse_partial_scores_env + + monkeypatch.setenv(PARTIAL_SCORES_ENV_VAR, "0") + assert _parse_partial_scores_env() is None + + +def test_parse_partial_scores_env_negative_disables(monkeypatch): + """Negative values disable partial scoring (returns None).""" + from tensorrt_llm.evaluate.lm_eval import PARTIAL_SCORES_ENV_VAR, _parse_partial_scores_env + + monkeypatch.setenv(PARTIAL_SCORES_ENV_VAR, "-5") + assert _parse_partial_scores_env() is None + + +def test_parse_partial_scores_env_invalid_raises(monkeypatch): + """A non-integer value raises ValueError.""" + from tensorrt_llm.evaluate.lm_eval import PARTIAL_SCORES_ENV_VAR, _parse_partial_scores_env + + monkeypatch.setenv(PARTIAL_SCORES_ENV_VAR, "abc") + with pytest.raises(ValueError, match=PARTIAL_SCORES_ENV_VAR): + _parse_partial_scores_env() + + +def test_parse_partial_scores_env_unset_returns_none(monkeypatch): + """Unset env var returns None.""" + from tensorrt_llm.evaluate.lm_eval import PARTIAL_SCORES_ENV_VAR, _parse_partial_scores_env + + monkeypatch.delenv(PARTIAL_SCORES_ENV_VAR, raising=False) + assert _parse_partial_scores_env() is None + + +# =========================================================================== +# LmEvalWrapper.generate_until — partial scorer invocation +# =========================================================================== + + +def test_generate_until_invokes_partial_scorer(): + """generate_until calls scorer.update and scorer.maybe_log for each response.""" + from tensorrt_llm.evaluate.lm_eval import LmEvalWrapper, _RunningScoreTracker + + fake_output = MagicMock() + fake_output.result.return_value.outputs = [MagicMock(text="42")] + fake_llm = MagicMock() + fake_llm.generate_async.return_value = fake_output + + wrapper = LmEvalWrapper( + llm=fake_llm, + partial_scores_every=1, + partial_scoring_task_dict={"fake_task": _FakeTask()}, + ) + + fake_request = MagicMock() + fake_request.args = ("hello world", {}) + fake_request.task_name = "fake_task" + fake_request.doc = {"answer": "42"} + + with ( + patch.object(_RunningScoreTracker, "update") as mock_update, + patch.object(_RunningScoreTracker, "maybe_log") as mock_log, + ): + wrapper.generate_until([fake_request], disable_tqdm=True) + + mock_update.assert_called_once() + mock_log.assert_called_once_with(1, 1) From fbff12dd21b1dddf68192454463c082f7efabf9e Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 24 Jul 2026 09:04:18 -0700 Subject: [PATCH 3/8] [None][feat] trtllm-eval: env-gated submission windowing for early partial-score signal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add TLLM_EVAL_MAX_IN_FLIGHT (default 0/unset = disabled, preserving existing submit-all behavior). When set to W>0, LmEvalWrapper.generate_until submits at most W requests concurrently; as each completes (via per-request waiter threads + FIRST_COMPLETED), it is scored through the existing _RunningScoreTracker and the window is topped up with the next unsubmitted request. Outputs are collected into an index-addressed list so the returned order still matches submission order — downstream scoring is unchanged. Why: on deployments using GUARANTEED_NO_EVICT scheduling the engine admits ~all requests concurrently, so every response completes in a burst at the very end and TLLM_EVAL_PARTIAL_SCORES_EVERY prints nothing until the final minute. Windowing makes responses complete steadily throughout the run, providing early failure signal. TRADEOFF: windowing can reduce end-to-end throughput — waves of W requests may under-fill the scheduler versus all-at-once admission. The payoff is EARLY FAILURE SIGNAL via steady partial scores. Suggested W: 256-512. Gated to the standard non-streaming path only; streaming and the multimodal generate_until override keep the submit-all behavior. A failed request re-raises without deadlocking the window (pool shutdown drains remaining waiters). Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 140 +++++++++++++++++++++++++++---- 1 file changed, 124 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 16003e77385d..fa53145d1bc5 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -12,6 +12,7 @@ # 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 concurrent.futures import copy import json import os @@ -55,6 +56,11 @@ # during long evals; see _RunningScoreTracker. 0/unset disables the feature. PARTIAL_SCORES_ENV_VAR = "TLLM_EVAL_PARTIAL_SCORES_EVERY" +# Cap on concurrently in-flight requests during LmEvalWrapper.generate_until. +# 0/unset disables windowing (submit everything up front — current behavior). +# See generate_until for the throughput/early-signal tradeoff. +MAX_IN_FLIGHT_ENV_VAR = "TLLM_EVAL_MAX_IN_FLIGHT" + class _RunningScoreTracker: """Best-effort running metric estimates over completed eval responses. @@ -177,6 +183,20 @@ def __init__(self, # signal for long evals. See _RunningScoreTracker. self.partial_scores_every = partial_scores_every self.partial_scoring_task_dict = partial_scoring_task_dict + # Env-gated cap on concurrently in-flight requests (0/unset = no cap, + # submit everything up front). See generate_until for the tradeoff. + env_window = os.environ.get(MAX_IN_FLIGHT_ENV_VAR) + if env_window: + try: + self.max_in_flight = int(env_window) + except ValueError: + raise ValueError( + f"{MAX_IN_FLIGHT_ENV_VAR} must be an integer, got " + f"{env_window!r}") from None + if self.max_in_flight < 0: + self.max_in_flight = 0 + else: + self.max_in_flight = 0 @property def eot_token_id(self) -> int: @@ -258,32 +278,117 @@ def _get_sampling_params(self, gen_kwargs: dict) -> SamplingParams: setattr(sampling_params, trtllm_key, value) return sampling_params - def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: - profiler.start("trtllm exec") - results = [] - for request in tqdm(requests, - desc="Submitting requests", - disable=disable_tqdm): - prompt, gen_kwargs = request.args + def _generate_until_windowed(self, requests, scorer, + disable_tqdm: bool) -> List[RequestOutput]: + """Submit requests through a sliding window of ``max_in_flight``. + + With unlimited submission and GUARANTEED_NO_EVICT scheduling on large + attention-DP deployments, the engine admits ~all requests concurrently + and every response completes in a burst at the very end, so the + partial-score tracker gives zero mid-run signal. Capping in-flight + requests to W makes responses complete steadily throughout the run: + as each request finishes, it is scored immediately and the window is + topped up with the next unsubmitted request. + + TRADEOFF: windowing can reduce end-to-end throughput — waves of W + requests may under-fill the scheduler compared with all-at-once + admission. The payoff is EARLY FAILURE SIGNAL via steady partial + scores, which is the right default for accuracy testing (kill a bad + run in minutes instead of hours). Suggested W: 256-512 for 16-rank + deployments. + + Completion (and therefore partial-scoring) order is arbitrary, but + outputs are collected into an index-addressed list so the returned + order matches submission order — downstream handling is unchanged. + """ + total = len(requests) + outputs: List[Optional[RequestOutput]] = [None] * total + next_idx = 0 + done_count = 0 + + def _submit_next(pool): + # generate_async stays on the caller thread (submission order is + # deterministic); only the blocking .result() wait is offloaded, + # one waiter thread per in-flight request, so completions surface + # in completion order via FIRST_COMPLETED below. + nonlocal next_idx + idx = next_idx + next_idx += 1 + prompt, gen_kwargs = requests[idx].args sampling_params = self._get_sampling_params(gen_kwargs) output = self.llm.generate_async(prompt, sampling_params=sampling_params, streaming=self.streaming) - results.append(output) + return pool.submit(lambda: (idx, output.result())) + + pbar = tqdm(total=total, desc="Fetching responses (windowed)", + disable=disable_tqdm) + try: + with concurrent.futures.ThreadPoolExecutor( + max_workers=self.max_in_flight) as pool: + pending = { + _submit_next(pool) + for _ in range(min(self.max_in_flight, total)) + } + while pending: + done, pending = concurrent.futures.wait( + pending, + return_when=concurrent.futures.FIRST_COMPLETED) + for fut in done: + # A failed request re-raises here (same as the + # non-windowed path's output.result()); no deadlock — + # pool shutdown just drains the remaining in-flight + # waiters, which the engine unblocks on completion or + # via EngineDeadError. + idx, output = fut.result() + outputs[idx] = output + done_count += 1 + pbar.update(1) + if scorer is not None: + scorer.update(requests[idx], + output.outputs[0].text) + scorer.maybe_log(done_count, total) + if next_idx < total: + pending.add(_submit_next(pool)) + finally: + pbar.close() + return outputs + + def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: + profiler.start("trtllm exec") scorer = None if self.partial_scores_every and self.partial_scoring_task_dict: scorer = _RunningScoreTracker(self.partial_scoring_task_dict, self.partial_scores_every) - outputs = [] - for output, request in zip( - tqdm(results, desc="Fetching responses", disable=disable_tqdm), - requests): - outputs.append(output.result()) - if scorer is not None: - scorer.update(request, outputs[-1].outputs[0].text) - scorer.maybe_log(len(outputs), len(requests)) + if self.max_in_flight > 0 and not self.streaming: + # Env-gated (TLLM_EVAL_MAX_IN_FLIGHT) submission windowing for the + # standard non-streaming path only; streaming (and the multimodal + # override below) keep the submit-all behavior untouched. + outputs = self._generate_until_windowed(requests, scorer, + disable_tqdm) + else: + results = [] + for request in tqdm(requests, + desc="Submitting requests", + disable=disable_tqdm): + prompt, gen_kwargs = request.args + sampling_params = self._get_sampling_params(gen_kwargs) + output = self.llm.generate_async( + prompt, + sampling_params=sampling_params, + streaming=self.streaming) + results.append(output) + + outputs = [] + for output, request in zip( + tqdm(results, desc="Fetching responses", + disable=disable_tqdm), requests): + outputs.append(output.result()) + if scorer is not None: + scorer.update(request, outputs[-1].outputs[0].text) + scorer.maybe_log(len(outputs), len(requests)) if self.output_dir: dump_inference_results(self.output_dir, outputs, @@ -519,6 +624,9 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: List of generated text responses """ profiler.start("trtllm exec") + # NOTE: TLLM_EVAL_MAX_IN_FLIGHT submission windowing (see + # LmEvalWrapper.generate_until) is intentionally NOT applied to this + # multimodal path; it keeps the original submit-all behavior. results = [] for request in tqdm(requests, desc="Submitting requests", From 103500fcdf536dc1a64427dd6c697ca3e59c3418 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 24 Jul 2026 09:39:34 -0700 Subject: [PATCH 4/8] [None][test] Strengthen generate_until scorer update assertion Assert exact (request, text) payload passed to _RunningScoreTracker.update instead of just checking call count. Addresses CodeRabbit review. Signed-off-by: Brian Nguyen --- tests/unittest/others/test_lm_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index 858df18f7e63..c318aeefab9a 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -948,5 +948,5 @@ def test_generate_until_invokes_partial_scorer(): ): wrapper.generate_until([fake_request], disable_tqdm=True) - mock_update.assert_called_once() + mock_update.assert_called_once_with(fake_request, "42") mock_log.assert_called_once_with(1, 1) From f748f982b120883683b05e2ecce65c969e08cebc Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 24 Jul 2026 12:59:15 -0500 Subject: [PATCH 5/8] Apply pre-commit auto-fixes Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index fa53145d1bc5..9c694b77d8f5 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -321,7 +321,8 @@ def _submit_next(pool): streaming=self.streaming) return pool.submit(lambda: (idx, output.result())) - pbar = tqdm(total=total, desc="Fetching responses (windowed)", + pbar = tqdm(total=total, + desc="Fetching responses (windowed)", disable=disable_tqdm) try: with concurrent.futures.ThreadPoolExecutor( @@ -332,8 +333,7 @@ def _submit_next(pool): } while pending: done, pending = concurrent.futures.wait( - pending, - return_when=concurrent.futures.FIRST_COMPLETED) + pending, return_when=concurrent.futures.FIRST_COMPLETED) for fut in done: # A failed request re-raises here (same as the # non-windowed path's output.result()); no deadlock — @@ -345,8 +345,7 @@ def _submit_next(pool): done_count += 1 pbar.update(1) if scorer is not None: - scorer.update(requests[idx], - output.outputs[0].text) + scorer.update(requests[idx], output.outputs[0].text) scorer.maybe_log(done_count, total) if next_idx < total: pending.add(_submit_next(pool)) @@ -383,7 +382,8 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: outputs = [] for output, request in zip( - tqdm(results, desc="Fetching responses", + tqdm(results, + desc="Fetching responses", disable=disable_tqdm), requests): outputs.append(output.result()) if scorer is not None: From 3ec1f1d2d5b5b8c8b55026b31ca7eb256354da46 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 27 Jul 2026 11:29:23 -0700 Subject: [PATCH 6/8] [None][feat] lm-eval: env-gated speculative acceptance-length summary Port TLLM_EVAL_SPEC_STATS from the Kimi K3 branch: when set to 1, generate_until logs a corpus-aggregate acceptance length (AL, mean of per-request avg_decoded_tokens_per_iter) after the run, for both the text and multimodal wrappers. Silent no-op on non-speculative runs. Unlike the original branch version, acceptance rate (AR) is intentionally not reported and return_perf_metrics is not forced: the request_perf_metrics.speculative_decoding counters AR needs are only maintained by TRTLLMSampler, not the TorchSampler used by one-engine spec-dec, so AR would silently read 0 on the default PyTorch path. AR reporting can return together with a fix populating those counters. Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 42 ++++++++ tests/unittest/others/test_lm_eval.py | 135 ++++++++++++++++++++++++++ 2 files changed, 177 insertions(+) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 9c694b77d8f5..99fe88970d05 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -61,6 +61,14 @@ # See generate_until for the throughput/early-signal tradeoff. MAX_IN_FLIGHT_ENV_VAR = "TLLM_EVAL_MAX_IN_FLIGHT" +# When "1", log an aggregate speculative-decoding summary (acceptance +# length AL as mean decoded tokens/step) at the end of generate_until. +# No-op output on non-speculative runs. Acceptance rate (AR) reporting is +# deferred: request_perf_metrics.speculative_decoding counters are only +# populated by TRTLLMSampler, not the TorchSampler used by one-engine +# spec-dec, so AR would silently read 0 on the default PyTorch path. +SPEC_STATS_ENV_VAR = "TLLM_EVAL_SPEC_STATS" + class _RunningScoreTracker: """Best-effort running metric estimates over completed eval responses. @@ -197,6 +205,9 @@ def __init__(self, self.max_in_flight = 0 else: self.max_in_flight = 0 + # Env-gated speculative-decoding stats (AL) aggregation over the + # eval corpus. See SPEC_STATS_ENV_VAR. + self.spec_stats = os.environ.get(SPEC_STATS_ENV_VAR) == "1" @property def eot_token_id(self) -> int: @@ -278,6 +289,31 @@ def _get_sampling_params(self, gen_kwargs: dict) -> SamplingParams: setattr(sampling_params, trtllm_key, value) return sampling_params + def _log_spec_stats(self, outputs: List[RequestOutput]) -> None: + """Log corpus-aggregate speculative-decoding stats (TLLM_EVAL_SPEC_STATS=1). + + AL (acceptance length) is reported as the mean of per-request + ``avg_decoded_tokens_per_iter`` (target token + accepted draft tokens + per decode step). Skips silently when the run produced no speculative + metrics (non-spec-dec config). + + Acceptance rate (AR) is intentionally not reported: the per-request + ``request_perf_metrics.speculative_decoding`` counters it needs are + only maintained by TRTLLMSampler (not the TorchSampler used by + one-engine spec-dec), so it would silently read 0 on the default + PyTorch path. Revisit once those counters are populated there. + """ + tokens_per_iter = [ + output.avg_decoded_tokens_per_iter for output in outputs + if getattr(output, "avg_decoded_tokens_per_iter", None) is not None + ] + if tokens_per_iter: + mean_tpi = sum(tokens_per_iter) / len(tokens_per_iter) + logger.info( + f"Spec-dec stats: AL (mean decoded tokens/step) {mean_tpi:.3f} " + f"(min {min(tokens_per_iter):.3f}, " + f"max {max(tokens_per_iter):.3f}, n={len(tokens_per_iter)})") + def _generate_until_windowed(self, requests, scorer, disable_tqdm: bool) -> List[RequestOutput]: """Submit requests through a sliding window of ``max_in_flight``. @@ -394,6 +430,9 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: dump_inference_results(self.output_dir, outputs, getattr(self.llm, 'tokenizer', None)) + if self.spec_stats: + self._log_spec_stats(outputs) + profiler.stop("trtllm exec") elapsed_time = profiler.elapsed_time_in_sec("trtllm exec") logger.info(f"TRTLLM execution time: {elapsed_time:.3f} seconds.") @@ -663,6 +702,9 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: dump_inference_results(self.output_dir, outputs, getattr(self.llm, 'tokenizer', None)) + if self.spec_stats: + self._log_spec_stats(outputs) + profiler.stop("trtllm exec") elapsed_time = profiler.elapsed_time_in_sec("trtllm exec") logger.info(f"TRTLLM execution time: {elapsed_time:.3f} seconds.") diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index c318aeefab9a..d4cf2cae6abc 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -21,6 +21,8 @@ * ``tensorrt_llm.evaluate.lm_eval_tasks.mmmu_pro.utils`` — ``parse_multi_choice_response`` reverse-scan and the ``MMMU_PRO_PROMPT_MODE`` env switch. +* ``LmEvalWrapper._log_spec_stats`` — the ``TLLM_EVAL_SPEC_STATS``-gated + speculative-decoding acceptance-length (AL) corpus summary. """ from __future__ import annotations @@ -950,3 +952,136 @@ def test_generate_until_invokes_partial_scorer(): mock_update.assert_called_once_with(fake_request, "42") mock_log.assert_called_once_with(1, 1) + + +# =========================================================================== +# TLLM_EVAL_SPEC_STATS — speculative-decoding AL stats +# =========================================================================== +# +# Only AL (acceptance length) is reported for now; AR needs the +# request_perf_metrics.speculative_decoding counters, which the TorchSampler +# used by one-engine spec-dec does not populate (see _log_spec_stats). + + +def _make_spec_output(tokens_per_iter: float | None = None) -> MagicMock: + """Fake RequestOutput with an optional per-request AL sample. + + ``tokens_per_iter`` as None models a request without speculative + metrics (non-spec-dec run, or a response that never reported them). + """ + output = MagicMock() + output.avg_decoded_tokens_per_iter = tokens_per_iter + output.outputs = [MagicMock()] + return output + + +def test_spec_stats_env_unset_disables(monkeypatch): + """Unset env leaves the feature off.""" + from tensorrt_llm.evaluate.lm_eval import SPEC_STATS_ENV_VAR + + monkeypatch.delenv(SPEC_STATS_ENV_VAR, raising=False) + wrapper = _make_lm_eval_wrapper() + assert wrapper.spec_stats is False + + +def test_spec_stats_env_enabled(monkeypatch): + """TLLM_EVAL_SPEC_STATS=1 turns the feature on.""" + from tensorrt_llm.evaluate.lm_eval import SPEC_STATS_ENV_VAR + + monkeypatch.setenv(SPEC_STATS_ENV_VAR, "1") + wrapper = _make_lm_eval_wrapper() + assert wrapper.spec_stats is True + + +@pytest.mark.parametrize("value", ["0", "true", "yes", ""]) +def test_spec_stats_env_non_one_values_disable(monkeypatch, value): + """Only the literal "1" enables the feature.""" + from tensorrt_llm.evaluate.lm_eval import SPEC_STATS_ENV_VAR + + monkeypatch.setenv(SPEC_STATS_ENV_VAR, value) + wrapper = _make_lm_eval_wrapper() + assert wrapper.spec_stats is False + + +def test_log_spec_stats_reports_al_mean_min_max(): + """AL is the mean of per-request avg_decoded_tokens_per_iter.""" + wrapper = _make_lm_eval_wrapper() + outputs = [ + _make_spec_output(tokens_per_iter=2.0), + _make_spec_output(tokens_per_iter=4.0), + ] + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + wrapper._log_spec_stats(outputs) + assert mock_logger.info.call_count == 1 + al_message = mock_logger.info.call_args[0][0] + assert "AL" in al_message + assert "3.000" in al_message # mean of 2.0 and 4.0 + assert "min 2.000" in al_message + assert "max 4.000" in al_message + assert "n=2" in al_message + + +def test_log_spec_stats_skips_requests_without_metrics(): + """Requests lacking spec metrics are excluded, not counted as zero.""" + wrapper = _make_lm_eval_wrapper() + outputs = [ + _make_spec_output(tokens_per_iter=3.0), + _make_spec_output(), # no metrics (e.g. dropped by the engine) + ] + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + wrapper._log_spec_stats(outputs) + assert mock_logger.info.call_count == 1 + message = mock_logger.info.call_args[0][0] + assert "3.000" in message + assert "n=1" in message + + +def test_log_spec_stats_silent_on_non_spec_run(): + """A run with no speculative metrics at all logs nothing.""" + wrapper = _make_lm_eval_wrapper() + outputs = [_make_spec_output(), _make_spec_output()] + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + wrapper._log_spec_stats(outputs) + mock_logger.info.assert_not_called() + + +def test_generate_until_logs_spec_stats_when_enabled(monkeypatch): + """generate_until forwards the collected outputs to _log_spec_stats.""" + from tensorrt_llm.evaluate.lm_eval import SPEC_STATS_ENV_VAR, LmEvalWrapper + + monkeypatch.setenv(SPEC_STATS_ENV_VAR, "1") + fake_result = MagicMock() + fake_result.outputs = [MagicMock(text="42")] + fake_output = MagicMock() + fake_output.result.return_value = fake_result + fake_llm = MagicMock() + fake_llm.generate_async.return_value = fake_output + + wrapper = LmEvalWrapper(llm=fake_llm) + fake_request = MagicMock() + fake_request.args = ("hello world", {}) + + with patch.object(LmEvalWrapper, "_log_spec_stats") as mock_stats: + wrapper.generate_until([fake_request], disable_tqdm=True) + + mock_stats.assert_called_once_with([fake_result]) + + +def test_generate_until_skips_spec_stats_when_disabled(monkeypatch): + """Without the env var, generate_until never touches _log_spec_stats.""" + from tensorrt_llm.evaluate.lm_eval import SPEC_STATS_ENV_VAR, LmEvalWrapper + + monkeypatch.delenv(SPEC_STATS_ENV_VAR, raising=False) + fake_output = MagicMock() + fake_output.result.return_value.outputs = [MagicMock(text="42")] + fake_llm = MagicMock() + fake_llm.generate_async.return_value = fake_output + + wrapper = LmEvalWrapper(llm=fake_llm) + fake_request = MagicMock() + fake_request.args = ("hello world", {}) + + with patch.object(LmEvalWrapper, "_log_spec_stats") as mock_stats: + wrapper.generate_until([fake_request], disable_tqdm=True) + + mock_stats.assert_not_called() From 794cd8c169a43ad80d873d659f91918de866e468 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 29 Jul 2026 21:57:14 -0700 Subject: [PATCH 7/8] [None][fix] Address review: CI wiring, weighted AL, windowed hardening - Add tests/unittest/others/test_lm_eval.py to l0_a10.yml so the suite actually runs in CI (test-db wires files individually; unlisted files are never collected). - Aggregate the TLLM_EVAL_SPEC_STATS acceptance length as an iteration-weighted mean (total decoded tokens / total decode iterations) using per-request decoding_iter, matching the canonical definition in bench/dataclasses/reporting.py; min/max stay per-request; missing decoding_iter falls back to weight 1. - Harden the TLLM_EVAL_MAX_IN_FLIGHT windowed path: bound waiter threads to min(W, len(requests)); on a failed request, escape via shutdown(wait=False, cancel_futures=True) instead of blocking on every other in-flight waiter (a hung request could previously trap the exception and hang the eval). Tests cover the raise path, the in-flight cap, and submission-order results under out-of-order completion. - Wire the partial scorer into the multimodal generate_until override (scoring the post-processed text lm-eval itself scores) instead of accepting and silently ignoring the arguments. - Document the shallow-copy / live-task caveat on _RunningScoreTracker. - Add end-to-end tests driving lm-eval's real evaluate() loop (real ConfigurableTask via custom_dataset, real filters and aggregation) over a mocked LLM: final score, partial-estimate-vs-final agreement, and windowed-path score parity. Signed-off-by: Brian Nguyen --- tensorrt_llm/evaluate/lm_eval.py | 156 +++--- .../integration/test_lists/test-db/l0_a10.yml | 1 + tests/unittest/others/test_lm_eval.py | 452 +++++++++++++++++- 3 files changed, 546 insertions(+), 63 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 99fe88970d05..f3128d605aae 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -86,6 +86,13 @@ class _RunningScoreTracker: by lm-eval is unaffected — scoring here happens on copies. Any failure (exotic task/filter/metric shapes) permanently disables the tracker for the run and never fails the eval itself. + + CAVEAT: the copy is shallow — ``probe.doc`` is shared with the live + instance, and filters / ``process_results`` run on the live task object. + lm-eval's filter contract is task/corpus-level and documents are mutable, + so a custom filter or a stateful ``process_results`` could in principle + touch state the final score depends on. The estimate is only meaningful + (and side-effect free) for stock per-document filter pipelines. """ def __init__(self, task_dict: dict, interval: int) -> None: @@ -292,10 +299,17 @@ def _get_sampling_params(self, gen_kwargs: dict) -> SamplingParams: def _log_spec_stats(self, outputs: List[RequestOutput]) -> None: """Log corpus-aggregate speculative-decoding stats (TLLM_EVAL_SPEC_STATS=1). - AL (acceptance length) is reported as the mean of per-request + AL (acceptance length) is the iteration-weighted mean of per-request ``avg_decoded_tokens_per_iter`` (target token + accepted draft tokens - per decode step). Skips silently when the run produced no speculative - metrics (non-spec-dec config). + per decode step), i.e. total decoded tokens / total decode iterations. + This matches the repo's canonical AL definition in + ``bench/dataclasses/reporting.py``: an equally-weighted mean would + bias the result toward short requests, which run fewer decode + iterations; iteration weighting makes it a token-level mean so longer + requests contribute proportionally. Requests that don't expose a + usable ``decoding_iter`` fall back to weight 1. min/max stay + per-request values. Skips silently when the run produced no + speculative metrics (non-spec-dec config). Acceptance rate (AR) is intentionally not reported: the per-request ``request_perf_metrics.speculative_decoding`` counters it needs are @@ -303,15 +317,23 @@ def _log_spec_stats(self, outputs: List[RequestOutput]) -> None: one-engine spec-dec), so it would silently read 0 on the default PyTorch path. Revisit once those counters are populated there. """ - tokens_per_iter = [ - output.avg_decoded_tokens_per_iter for output in outputs - if getattr(output, "avg_decoded_tokens_per_iter", None) is not None - ] - if tokens_per_iter: - mean_tpi = sum(tokens_per_iter) / len(tokens_per_iter) + samples = [] # (avg_decoded_tokens_per_iter, weight=decode iterations) + for output in outputs: + tpi = getattr(output, "avg_decoded_tokens_per_iter", None) + if tpi is None: + continue + iters = getattr(output, "decoding_iter", None) + weight = iters if isinstance(iters, (int, float)) and iters > 0 \ + else 1 + samples.append((tpi, weight)) + if samples: + weighted_al = (sum(tpi * w for tpi, w in samples) / + sum(w for _, w in samples)) + tokens_per_iter = [tpi for tpi, _ in samples] logger.info( - f"Spec-dec stats: AL (mean decoded tokens/step) {mean_tpi:.3f} " - f"(min {min(tokens_per_iter):.3f}, " + f"Spec-dec stats: AL (decoded tokens/step, " + f"iteration-weighted) {weighted_al:.3f} " + f"(per-request min {min(tokens_per_iter):.3f}, " f"max {max(tokens_per_iter):.3f}, n={len(tokens_per_iter)})") def _generate_until_windowed(self, requests, scorer, @@ -338,6 +360,8 @@ def _generate_until_windowed(self, requests, scorer, order matches submission order — downstream handling is unchanged. """ total = len(requests) + if total == 0: + return [] outputs: List[Optional[RequestOutput]] = [None] * total next_idx = 0 done_count = 0 @@ -360,42 +384,56 @@ def _submit_next(pool): pbar = tqdm(total=total, desc="Fetching responses (windowed)", disable=disable_tqdm) + # The window size also bounds the waiter-thread count, so cap it at + # the request count (W >= total would otherwise spawn one idle-capable + # thread per request). + pool = concurrent.futures.ThreadPoolExecutor( + max_workers=min(self.max_in_flight, total)) try: - with concurrent.futures.ThreadPoolExecutor( - max_workers=self.max_in_flight) as pool: - pending = { - _submit_next(pool) - for _ in range(min(self.max_in_flight, total)) - } - while pending: - done, pending = concurrent.futures.wait( - pending, return_when=concurrent.futures.FIRST_COMPLETED) - for fut in done: - # A failed request re-raises here (same as the - # non-windowed path's output.result()); no deadlock — - # pool shutdown just drains the remaining in-flight - # waiters, which the engine unblocks on completion or - # via EngineDeadError. - idx, output = fut.result() - outputs[idx] = output - done_count += 1 - pbar.update(1) - if scorer is not None: - scorer.update(requests[idx], output.outputs[0].text) - scorer.maybe_log(done_count, total) - if next_idx < total: - pending.add(_submit_next(pool)) + pending = { + _submit_next(pool) + for _ in range(min(self.max_in_flight, total)) + } + while pending: + done, pending = concurrent.futures.wait( + pending, return_when=concurrent.futures.FIRST_COMPLETED) + for fut in done: + # A failed request re-raises here (same as the + # non-windowed path's output.result()). + idx, output = fut.result() + outputs[idx] = output + done_count += 1 + pbar.update(1) + if scorer is not None: + scorer.update(requests[idx], output.outputs[0].text) + scorer.maybe_log(done_count, total) + if next_idx < total: + pending.add(_submit_next(pool)) + except BaseException: + # Fail fast: a blocking shutdown here (as a `with` block's + # __exit__ would do) joins every other in-flight waiter with no + # cancellation or timeout — if any of those never resolves, the + # exception can't escape and the eval hangs instead of failing. + # shutdown(wait=False) lets the exception propagate immediately; + # already-running waiters unblock when the engine completes or + # kills their requests (e.g. EngineDeadError). + pool.shutdown(wait=False, cancel_futures=True) + raise finally: pbar.close() + pool.shutdown(wait=True) return outputs + def _make_partial_scorer(self) -> Optional["_RunningScoreTracker"]: + if self.partial_scores_every and self.partial_scoring_task_dict: + return _RunningScoreTracker(self.partial_scoring_task_dict, + self.partial_scores_every) + return None + def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: profiler.start("trtllm exec") - scorer = None - if self.partial_scores_every and self.partial_scoring_task_dict: - scorer = _RunningScoreTracker(self.partial_scoring_task_dict, - self.partial_scores_every) + scorer = self._make_partial_scorer() if self.max_in_flight > 0 and not self.streaming: # Env-gated (TLLM_EVAL_MAX_IN_FLIGHT) submission windowing for the @@ -481,9 +519,10 @@ def __init__(self, to model outputs before scoring. Used by Kimi K2.5 to strip ``...`` and extract the final answer (see ``tensorrt_llm.evaluate.post_processing``). - partial_scores_every: Accepted for interface parity with - LmEvalWrapper; the multimodal generate_until override does - not implement partial scoring yet. + partial_scores_every: Same as LmEvalWrapper — log running + metric estimates every N completed responses. The multimodal + path scores the post-processed text (the same string lm-eval + itself scores). partial_scoring_task_dict: See partial_scores_every. """ super().__init__( @@ -666,6 +705,7 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: # NOTE: TLLM_EVAL_MAX_IN_FLIGHT submission windowing (see # LmEvalWrapper.generate_until) is intentionally NOT applied to this # multimodal path; it keeps the original submit-all behavior. + scorer = self._make_partial_scorer() results = [] for request in tqdm(requests, desc="Submitting requests", @@ -693,10 +733,24 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: results.append(output) outputs = [] - for output in tqdm(results, - desc="Fetching responses", - disable=disable_tqdm): + results_text = [] + for output, request in zip( + tqdm(results, desc="Fetching responses", disable=disable_tqdm), + requests): outputs.append(output.result()) + # Apply per-sample post-processing only when caller injected one. + # Kimi K2.5 passes strip_thinking_and_extract_mmmu_answer to + # recover answers from ...-wrapped outputs that + # lm-eval's default extractor cannot parse. + raw = outputs[-1].outputs[0].text + text = self.post_process_fn( + raw) if self.post_process_fn is not None else raw + results_text.append(text) + if scorer is not None: + # Score the post-processed text — the same string lm-eval + # itself will score. + scorer.update(request, text) + scorer.maybe_log(len(outputs), len(requests)) if self.output_dir: dump_inference_results(self.output_dir, outputs, @@ -710,18 +764,6 @@ def generate_until(self, requests, disable_tqdm: bool = False) -> List[str]: logger.info(f"TRTLLM execution time: {elapsed_time:.3f} seconds.") profiler.reset("trtllm exec") - # Apply per-sample post-processing only when caller injected one. - # Kimi K2.5 passes strip_thinking_and_extract_mmmu_answer to recover - # answers from ...-wrapped outputs that lm-eval's - # default extractor cannot parse. - results_text = [] - for output in outputs: - raw = output.outputs[0].text - if self.post_process_fn is not None: - results_text.append(self.post_process_fn(raw)) - else: - results_text.append(raw) - return results_text diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index e4bcd807415d..c36a94225c63 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -59,6 +59,7 @@ l0_a10: - unittest/others/test_cache_transceiver_precheck_config.py - unittest/others/test_cache_transceiver_precheck_run.py - unittest/others/test_convert_utils.py + - unittest/others/test_lm_eval.py - unittest/others/test_lora_manager.py - unittest/others/test_lora_module_count.py - unittest/others/test_time_breakdown.py diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index d4cf2cae6abc..9e1fcc02fc0d 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -22,13 +22,23 @@ ``parse_multi_choice_response`` reverse-scan and the ``MMMU_PRO_PROMPT_MODE`` env switch. * ``LmEvalWrapper._log_spec_stats`` — the ``TLLM_EVAL_SPEC_STATS``-gated - speculative-decoding acceptance-length (AL) corpus summary. + speculative-decoding acceptance-length (AL) corpus summary, + iteration-weighted to match ``bench/dataclasses/reporting.py``. +* ``LmEvalWrapper._generate_until_windowed`` — the + ``TLLM_EVAL_MAX_IN_FLIGHT`` submission window: in-flight cap, + submission-order results under out-of-order completion, and fail-fast + propagation of request errors. +* End-to-end: lm-eval's real ``evaluate()`` loop (real ``ConfigurableTask``, + filters, aggregation) driven through ``LmEvalWrapper`` over a mocked LLM, + for both the final score and the partial-score running estimates. """ from __future__ import annotations import importlib import os +import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -36,6 +46,7 @@ from tensorrt_llm.evaluate.covost2 import CoVoST2 from tensorrt_llm.evaluate.lm_eval import ( LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER, + MAX_IN_FLIGHT_ENV_VAR, LmEvalWrapper, MultimodalLmEvalWrapper, ) @@ -961,16 +972,25 @@ def test_generate_until_invokes_partial_scorer(): # Only AL (acceptance length) is reported for now; AR needs the # request_perf_metrics.speculative_decoding counters, which the TorchSampler # used by one-engine spec-dec does not populate (see _log_spec_stats). +# AL is iteration-weighted (total decoded tokens / total decode iterations) +# to agree with the repo's canonical definition in +# ``bench/dataclasses/reporting.py``. -def _make_spec_output(tokens_per_iter: float | None = None) -> MagicMock: +def _make_spec_output( + tokens_per_iter: float | None = None, + decoding_iter: int | None = 1, +) -> MagicMock: """Fake RequestOutput with an optional per-request AL sample. ``tokens_per_iter`` as None models a request without speculative metrics (non-spec-dec run, or a response that never reported them). + ``decoding_iter`` is the AL aggregation weight (decode iterations the + request ran); None models a result that never populated it. """ output = MagicMock() output.avg_decoded_tokens_per_iter = tokens_per_iter + output.decoding_iter = decoding_iter output.outputs = [MagicMock()] return output @@ -1004,23 +1024,58 @@ def test_spec_stats_env_non_one_values_disable(monkeypatch, value): def test_log_spec_stats_reports_al_mean_min_max(): - """AL is the mean of per-request avg_decoded_tokens_per_iter.""" + """AL over equal-weight requests equals the plain mean; min/max/n present.""" wrapper = _make_lm_eval_wrapper() outputs = [ - _make_spec_output(tokens_per_iter=2.0), - _make_spec_output(tokens_per_iter=4.0), + _make_spec_output(tokens_per_iter=2.0, decoding_iter=5), + _make_spec_output(tokens_per_iter=4.0, decoding_iter=5), ] with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: wrapper._log_spec_stats(outputs) assert mock_logger.info.call_count == 1 al_message = mock_logger.info.call_args[0][0] assert "AL" in al_message - assert "3.000" in al_message # mean of 2.0 and 4.0 + assert "3.000" in al_message # equal weights -> mean of 2.0 and 4.0 assert "min 2.000" in al_message assert "max 4.000" in al_message assert "n=2" in al_message +def test_log_spec_stats_weights_by_decode_iterations(): + """AL matches reporting.py's token-level mean: weighted by decode iterations. + + (2.0 tok/iter over 1 iter) + (4.0 tok/iter over 3 iters) = 14 decoded + tokens over 4 iterations = 3.5 — NOT the unweighted mean 3.0, which + would bias the result toward short requests (see the explicit rationale + in ``bench/dataclasses/reporting.py``). + """ + wrapper = _make_lm_eval_wrapper() + outputs = [ + _make_spec_output(tokens_per_iter=2.0, decoding_iter=1), + _make_spec_output(tokens_per_iter=4.0, decoding_iter=3), + ] + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + wrapper._log_spec_stats(outputs) + message = mock_logger.info.call_args[0][0] + assert "3.500" in message + # min/max stay per-request values, unweighted. + assert "min 2.000" in message + assert "max 4.000" in message + + +def test_log_spec_stats_missing_decoding_iter_falls_back_to_weight_one(): + """Requests without a usable decoding_iter contribute with weight 1.""" + wrapper = _make_lm_eval_wrapper() + outputs = [ + _make_spec_output(tokens_per_iter=2.0, decoding_iter=None), + _make_spec_output(tokens_per_iter=4.0, decoding_iter=0), + ] + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + wrapper._log_spec_stats(outputs) + message = mock_logger.info.call_args[0][0] + assert "3.000" in message # both fall back to weight 1 -> plain mean + + def test_log_spec_stats_skips_requests_without_metrics(): """Requests lacking spec metrics are excluded, not counted as zero.""" wrapper = _make_lm_eval_wrapper() @@ -1085,3 +1140,388 @@ def test_generate_until_skips_spec_stats_when_disabled(monkeypatch): wrapper.generate_until([fake_request], disable_tqdm=True) mock_stats.assert_not_called() + + +# =========================================================================== +# TLLM_EVAL_MAX_IN_FLIGHT — windowed generate_until +# =========================================================================== +# +# The windowed path caps concurrently in-flight requests at W, tops the +# window up as responses complete, and collects outputs into an +# index-addressed list. The correctness property the whole design exists to +# preserve is SUBMISSION-ORDER RESULTS under arbitrary completion order; +# the liveness property is that a failed request propagates promptly +# instead of deadlocking behind other in-flight waiters. + + +class _FakeAsyncOutput: + """Async handle whose blocking .result() is supplied by the test.""" + + def __init__(self, result_fn): + self._result_fn = result_fn + + def result(self): + return self._result_fn() + + +def _text_result(text: str) -> MagicMock: + result = MagicMock() + result.outputs = [MagicMock(text=text)] + return result + + +def _make_windowed_llm(events, error_idx=None): + """LLM whose request i blocks until events[i] is set, then yields resp-i. + + Returns the fake llm and the (mutated) list of submitted request indices, + so tests can observe how far submission has progressed. + """ + submitted = [] + llm = MagicMock() + llm.tokenizer = MagicMock() + + def generate_async(prompt, sampling_params=None, streaming=False): + idx = len(submitted) + submitted.append(idx) + + def _result(): + assert events[idx].wait(timeout=30), f"request {idx} never released" + if error_idx is not None and idx == error_idx: + raise RuntimeError(f"request {idx} failed") + return _text_result(f"resp-{idx}") + + return _FakeAsyncOutput(_result) + + llm.generate_async = generate_async + return llm, submitted + + +def _make_requests(n: int) -> list: + requests = [] + for i in range(n): + request = MagicMock() + request.args = (f"prompt-{i}", {}) + requests.append(request) + return requests + + +def _wait_until(predicate, timeout: float = 10.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +def test_windowed_caps_in_flight_and_preserves_order(monkeypatch): + """At most W requests in flight, and results follow submission order. + + Request 1 is completed before request 0 (out-of-order completion); the + window tops up with request 2 only after that completion, and the + returned list is still resp-0..resp-4 in submission order. + """ + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "2") + total = 5 + events = [threading.Event() for _ in range(total)] + llm, submitted = _make_windowed_llm(events) + wrapper = LmEvalWrapper(llm=llm) + assert wrapper.max_in_flight == 2 + + returned = [] + worker = threading.Thread( + target=lambda: returned.append( + wrapper.generate_until(_make_requests(total), disable_tqdm=True) + ) + ) + worker.start() + try: + # Only the first W requests are submitted while none have completed. + assert _wait_until(lambda: len(submitted) == 2) + time.sleep(0.05) + assert len(submitted) == 2, "window overshot max_in_flight" + # Completing request 1 (out of order) tops the window up by one. + events[1].set() + assert _wait_until(lambda: len(submitted) == 3) + time.sleep(0.05) + assert len(submitted) == 3 + finally: + for event in events: + event.set() + worker.join(timeout=30) + assert not worker.is_alive() + assert returned and returned[0] == [f"resp-{i}" for i in range(total)] + + +def test_windowed_failed_request_raises_without_waiting(monkeypatch): + """A failed request propagates while another request is still in flight. + + Regression guard for the deadlock the review called out: a blocking + pool shutdown would join every other outstanding waiter with no + cancellation or timeout, so if any of them never resolved the exception + could not escape and the eval hung instead of failing. + """ + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "2") + events = [threading.Event() for _ in range(2)] + llm, _ = _make_windowed_llm(events, error_idx=0) + wrapper = LmEvalWrapper(llm=llm) + events[0].set() # request 0 fails immediately; request 1 stays blocked + try: + with pytest.raises(RuntimeError, match="request 0 failed"): + wrapper.generate_until(_make_requests(2), disable_tqdm=True) + # The exception escaped while request 1 had not resolved. + assert not events[1].is_set() + finally: + events[1].set() # release the lingering waiter thread + + +def test_windowed_window_larger_than_request_count(monkeypatch): + """W >= len(requests) submits each request exactly once and stays ordered.""" + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "64") + total = 3 + events = [threading.Event() for _ in range(total)] + for event in events: + event.set() + llm, submitted = _make_windowed_llm(events) + wrapper = LmEvalWrapper(llm=llm) + result = wrapper.generate_until(_make_requests(total), disable_tqdm=True) + assert result == [f"resp-{i}" for i in range(total)] + assert submitted == list(range(total)) + + +def test_windowed_empty_request_list(monkeypatch): + """Zero requests short-circuit without creating a thread pool.""" + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "2") + llm, _ = _make_windowed_llm([]) + wrapper = LmEvalWrapper(llm=llm) + assert wrapper.generate_until([], disable_tqdm=True) == [] + + +def test_windowed_invokes_partial_scorer(monkeypatch): + """The windowed path feeds every completion to the partial scorer.""" + from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "2") + total = 3 + events = [threading.Event() for _ in range(total)] + for event in events: + event.set() + llm, _ = _make_windowed_llm(events) + wrapper = LmEvalWrapper( + llm=llm, + partial_scores_every=1, + partial_scoring_task_dict={"fake_task": _FakeTask()}, + ) + with ( + patch.object(_RunningScoreTracker, "update") as mock_update, + patch.object(_RunningScoreTracker, "maybe_log") as mock_log, + ): + wrapper.generate_until(_make_requests(total), disable_tqdm=True) + assert mock_update.call_count == total + assert mock_log.call_count == total + + +# =========================================================================== +# MultimodalLmEvalWrapper.generate_until — partial scorer wiring +# =========================================================================== + + +def test_multimodal_generate_until_invokes_partial_scorer(): + """The multimodal override scores the post-processed text lm-eval sees.""" + from tensorrt_llm.evaluate.lm_eval import _RunningScoreTracker + + fake_output = MagicMock() + fake_output.result.return_value.outputs = [MagicMock(text="reasoning42")] + fake_llm = MagicMock() + fake_llm.tokenizer = MagicMock() + fake_llm.input_processor = MagicMock() + fake_llm.generate_async.return_value = fake_output + + with patch.object(MultimodalLmEvalWrapper, "_get_model_type", return_value="gemma3"): + wrapper = MultimodalLmEvalWrapper( + fake_llm, + sampling_params=None, + model_type="gemma3", + post_process_fn=lambda s: s.split("")[-1], + partial_scores_every=1, + partial_scoring_task_dict={"fake_task": _FakeTask()}, + ) + + fake_request = MagicMock() + fake_request.args = ("prompt", {}, {"visual": [MagicMock()]}) + + with ( + patch( + "tensorrt_llm.evaluate.lm_eval.prompt_inputs", + side_effect=lambda p: {"prompt": p}, + ), + patch( + "tensorrt_llm.evaluate.lm_eval.convert_image_mode", + side_effect=lambda img, mode: img, + ), + patch.object(_RunningScoreTracker, "update") as mock_update, + patch.object(_RunningScoreTracker, "maybe_log") as mock_log, + ): + results = wrapper.generate_until([fake_request], disable_tqdm=True) + + assert results == ["42"] + # The scorer must see the post-processed text, not the raw output. + mock_update.assert_called_once_with(fake_request, "42") + mock_log.assert_called_once_with(1, 1) + + +# =========================================================================== +# End-to-end: real lm-eval evaluator over a mocked LLM +# =========================================================================== +# +# Runs lm_eval.evaluator.evaluate() — real ConfigurableTask, real filter +# pipeline, real metric aggregation — against LmEvalWrapper wrapping a +# mocked LLM that returns canned responses. This exercises the exact +# calling conventions between the harness and the wrapper (instance +# shapes, filtered_resps, the process_results list convention) that pure +# unit mocks can get subtly wrong; the GSM8K first-character bug above +# survived precisely because nothing ran the real harness loop. + + +_E2E_DOCS = [ + {"question": "2+2?", "answer": "4"}, + {"question": "3+4?", "answer": "7"}, + {"question": "5+6?", "answer": "11"}, + {"question": "10-3?", "answer": "7"}, +] +# Model answers: 3 correct, 1 wrong ("12" != "11") -> exact_match 0.75. +_E2E_RESPONSES = [ + "The answer is 4.", + "The answer is 7.", + "The answer is 12.", + "The answer is 7.", +] + + +def _toy_task(): + """A real generate_until ConfigurableTask over an in-memory dataset.""" + import datasets + from lm_eval.api.task import ConfigurableTask + + return ConfigurableTask( + config={ + "task": "toy_arith", + "custom_dataset": lambda **kwargs: datasets.DatasetDict( + {"test": datasets.Dataset.from_list(_E2E_DOCS)} + ), + "test_split": "test", + "output_type": "generate_until", + "doc_to_text": "Q: {{question}}\nA:", + "doc_to_target": "{{answer}}", + "generation_kwargs": {"until": ["\n"], "do_sample": False}, + "filter_list": [ + { + "name": "strict-match", + "filter": [ + {"function": "regex", "regex_pattern": r"(-?[0-9]+)"}, + {"function": "take_first"}, + ], + } + ], + "metric_list": [ + { + "metric": "exact_match", + "aggregation": "mean", + "higher_is_better": True, + } + ], + } + ) + + +def _canned_llm(responses): + """LLM whose generate_async yields the canned texts in submission order.""" + llm = MagicMock() + llm.tokenizer = MagicMock() + response_iter = iter(responses) + + def generate_async(prompt, sampling_params=None, streaming=False): + text = next(response_iter) + output = MagicMock() + output.result.return_value = _text_result(text) + return output + + llm.generate_async = generate_async + return llm + + +def test_e2e_harness_final_score_over_mocked_llm(): + """The real lm-eval evaluator scores canned responses correctly.""" + from lm_eval.evaluator import evaluate + + task = _toy_task() + wrapper = LmEvalWrapper(llm=_canned_llm(_E2E_RESPONSES)) + results = evaluate( + lm=wrapper, + task_dict={"toy_arith": task}, + bootstrap_iters=0, + log_samples=False, + ) + score = results["results"]["toy_arith"]["exact_match,strict-match"] + assert score == pytest.approx(0.75) + + +def test_e2e_partial_scores_match_final_score(): + """Partial-score estimates over the full corpus agree with the harness. + + Uses the REAL task's filters and process_results inside + _RunningScoreTracker (no fakes), so a calling-convention mismatch + between the tracker and lm-eval internals disables the tracker and + fails this test. + """ + from lm_eval.evaluator import evaluate + + task = _toy_task() + task_dict = {"toy_arith": task} + wrapper = LmEvalWrapper( + llm=_canned_llm(_E2E_RESPONSES), + partial_scores_every=2, + partial_scoring_task_dict=task_dict, + ) + with patch("tensorrt_llm.evaluate.lm_eval.logger") as mock_logger: + results = evaluate( + lm=wrapper, + task_dict=task_dict, + bootstrap_iters=0, + log_samples=False, + ) + messages = [call.args[0] for call in mock_logger.info.call_args_list] + assert not any("Partial scoring disabled" in m for m in messages), ( + "tracker was disabled by a scoring failure against the real task" + ) + partial = [m for m in messages if "Partial scores" in m] + # interval=2 over 4 responses -> logs at 2/4 and 4/4. + assert len(partial) == 2 + assert "2/4" in partial[0] + assert "4/4" in partial[1] + # The final running estimate agrees with the true score (0~100 scale). + assert "75.00" in partial[1] + score = results["results"]["toy_arith"]["exact_match,strict-match"] + assert score == pytest.approx(0.75) + + +def test_e2e_windowed_matches_final_score(monkeypatch): + """The windowed path produces the same harness score as submit-all. + + Windowing must be a pure scheduling change — outputs are collected in + submission order, so the score is identical to the default path. + """ + from lm_eval.evaluator import evaluate + + monkeypatch.setenv(MAX_IN_FLIGHT_ENV_VAR, "2") + task = _toy_task() + wrapper = LmEvalWrapper(llm=_canned_llm(_E2E_RESPONSES)) + assert wrapper.max_in_flight == 2 + results = evaluate( + lm=wrapper, + task_dict={"toy_arith": task}, + bootstrap_iters=0, + log_samples=False, + ) + score = results["results"]["toy_arith"]["exact_match,strict-match"] + assert score == pytest.approx(0.75) From 2d4c3913e8185263f59a3c683ccb62082dede064 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 30 Jul 2026 05:21:15 -0700 Subject: [PATCH 8/8] [None][test] Fix rotted multi-image interleave tests before L0 wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three multi-image content_parts tests constructed the wrapper with model_type="gemma3", which never opts into interleaved placeholders (MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders returns False for it), so apply_chat_template never built content_parts and the tests failed with KeyError — silently, since this file was not collected by any CI list. Now that l0_a10.yml collects the file, they must be green: drive the registry flag explicitly (interleave=True/False) so the tests are independent of which models the installed transformers version registers, and rename the two negative-path tests to describe what they actually assert. Full file passes 87/87. Signed-off-by: Brian Nguyen --- tests/unittest/others/test_lm_eval.py | 39 +++++++++++++++++++-------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index 9e1fcc02fc0d..a7d20a97a8f1 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -58,6 +58,7 @@ strip_string, ) from tensorrt_llm.inputs.content_format import ContentFormat +from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.sampling_params import SamplingParams # =========================================================================== @@ -301,14 +302,30 @@ def test_sampling_override_no_cli_falls_back_to_yaml(): # correctly-ordered OpenAI content list. -# Uses ``gemma3`` by default because it is always registered regardless of -# transformers version; the wrapper's interleave logic itself is generic. -def _make_multimodal_wrapper(model_type: str = "gemma3") -> MultimodalLmEvalWrapper: +# Interleaving is opt-in per model: the wrapper reads +# ``MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders(model_type)`` +# at construction, and models that don't opt in keep the historical +# strip-and-bulk-insert behaviour. These tests drive that flag directly +# instead of naming an opted-in model, because which models are registered +# varies with the installed transformers version — keying on a real model +# name would make the tests environment-dependent. ``interleave=False`` +# (the default here) matches an unregistered model such as ``gemma3``. +def _make_multimodal_wrapper( + model_type: str = "gemma3", + interleave: bool = False, +) -> MultimodalLmEvalWrapper: fake_llm = MagicMock() fake_llm.tokenizer = MagicMock() fake_llm.input_processor = MagicMock() fake_llm.input_processor.processor = MagicMock() - with patch.object(MultimodalLmEvalWrapper, "_get_model_type", return_value=model_type): + with ( + patch.object(MultimodalLmEvalWrapper, "_get_model_type", return_value=model_type), + patch.object( + MULTIMODAL_PLACEHOLDER_REGISTRY, + "get_interleave_placeholders", + return_value=interleave, + ), + ): return MultimodalLmEvalWrapper( fake_llm, sampling_params=None, @@ -350,8 +367,8 @@ def _fake_trtllm_apply(**kwargs): return convs[0] -def test_single_image_does_not_interleave(): - """Single-image prompts never need interleaving. +def test_not_opted_in_model_does_not_interleave(): + """A model that does not opt in keeps the historical bulk-insert path. content_parts stays absent so the existing BEFORE_TEXT default keeps working. """ @@ -367,7 +384,7 @@ def test_multi_image_openai_builds_content_parts(): ``_build_openai_content`` then emits media entries at the correct positions. """ - wrapper = _make_multimodal_wrapper() + wrapper = _make_multimodal_wrapper(interleave=True) ph = LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER text = f"Consider {ph}. What does {ph} show?" conv = _call_apply(wrapper, text, content_format=ContentFormat.OPENAI) @@ -383,8 +400,8 @@ def test_multi_image_openai_builds_content_parts(): assert [p["media_index"] for p in media_parts] == [0, 1] -def test_multi_image_string_format_skips_interleave(): - """STRING-format chat templates skip the interleaving path. +def test_multi_image_string_format_not_opted_in_uses_placeholders(): + """STRING-format templates on a non-opted-in model use flat placeholders. Placeholders are inserted into the flat text via ``add_multimodal_placeholders`` instead, so ``content_parts`` stays absent. @@ -406,7 +423,7 @@ def test_trailing_text_after_last_image_preserved(): Otherwise the question suffix ('Answer:') is dropped before it reaches the model. """ - wrapper = _make_multimodal_wrapper() + wrapper = _make_multimodal_wrapper(interleave=True) ph = LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER text = f"Compare {ph} with {ph}. Answer with a letter." conv = _call_apply(wrapper, text, content_format=ContentFormat.OPENAI) @@ -421,7 +438,7 @@ def test_leading_image_no_empty_text_segment(): content_parts must begin with the image entry itself. """ - wrapper = _make_multimodal_wrapper() + wrapper = _make_multimodal_wrapper(interleave=True) ph = LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER text = f"{ph} {ph} Answer?" conv = _call_apply(wrapper, text, content_format=ContentFormat.OPENAI)