[#17022][fix] trtllm-eval: make --max_output_length reachable on aime25/aime26

[thorjohnsen]

Description

--max_output_length is advertised as "Maximum generation length" on every trtllm-eval subcommand. It is placed into SamplingParams.max_tokens (lm_eval.py:715, :1196), and then _get_sampling_params unconditionally overwrites it with the task yaml's max_gen_toks. Passing a larger value does nothing, silently, and there is no workaround on text tasks.

The guard that exists for exactly this case, preserve_caller_max_tokens, was unreachable from any text task because of two independent gates:

1. Not forwarded. evaluate() passed it to the wrapper only for multimodal (lm_eval.py:623-627), so the text LmEvalWrapper kept its default False (lm_eval.py:65).
2. Not exposed. --preserve_caller_max_tokens was declared on exactly one subcommand, mmmu (lm_eval.py:1269).

The guard logic itself was already correct — just unreachable.

Note the asymmetry that makes this surprising: --max_input_length maps to truncate_prompt_tokens, which is not in params_mapping, so it works as documented. --temperature / --top_p have a dedicated escape hatch (sampling_override). Only the output budget is clobbered with no way out.

Which tasks are affected depends on whether the yaml sets max_gen_toks: aime25 / aime26 pin 32768; gsm8k and gpqa_diamond leave it unset, so the CLI value already survives there.

The change

• Forward preserve_caller_max_tokens to both wrappers. LmEvalWrapper already accepts the argument; only the forwarding was conditional.
• Expose --preserve_caller_max_tokens on aime25 and aime26.
• Warn once per run when a caller-supplied max_tokens is discarded in favor of the yaml value. The clamp is usually the right default — matching the reference harness recipe — so it stays on. The hazard was that it was invisible.

Default behavior is unchanged: without the flag, the yaml still wins.

This deliberately does not touch longbench_v1, which #17022 also listed. On closer reading that subcommand exposes neither --max_output_length nor --max_input_length and passes sampling_params=None, so there is no caller value to preserve and the flag would be inert. Raising its per-subtask caps (32–512, appropriate for short extractive answers but too small for a thinking model) would be a separate change.

Reported as #17022.

Test Coverage

No unit test is added — _get_sampling_params needs an LLM instance to construct a wrapper through the normal path, and the existing tests/unittest/llmapi/ suite has no lm-eval wrapper harness to extend. Suggestions on where this would best live are welcome.

Verified directly instead, on DeepSeek-V4-Flash (4xGB300, TP4/EP4, greedy):

The defect. On aime25 at the 32768 cap, 3 of 30 problems (docs 13, 19, 23) ran past it mid-chain-of-thought, produced 91k–102k characters, never emitted </think>, and scored 0. --max_output_length 131072 had no effect — the yaml's 32768 was reapplied per request.

The fix. The equivalent of this patch, applied as an in-process monkeypatch forcing preserve_caller_max_tokens=True on the text wrapper, with _get_sampling_params instrumented to prove it took:

[PATCH] preserve_caller_max_tokens=True -> resolved max_tokens=131072
        (task yaml max_gen_toks would be 32768)

2 of those 3 problems then terminated cleanly and scored correct (86.67 → 93.33 exact_match; single greedy run each, so per-run variance applies — see the discussion on #17022).

Behavioral checks against this branch, exercising _get_sampling_params directly:

caller max_tokens	yaml max_gen_toks	flag	resolved
131072	32768	off	32768 (unchanged default, warns once)
131072	32768	on	131072
1024	32768	on	32768 (guard only keeps the larger caller value)

Also confirmed the warning fires exactly once across repeated requests, and that --preserve_caller_max_tokens is now registered on aime25 / aime26 / mmmu and absent from gsm8k / longbench_v1.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.