feat(speculative): add multi-dataset acceptance-length benchmark sweep#2966
Conversation
khazic
requested review from
a team,
HuiyingLi,
akoumpa,
athitten and
snowmanwwg
as code owners
bench_sglang.py / bench_vllm.py each measure one dataset per invocation, but the same draft's acceptance rate can vary sharply by task -- conversational data tends to accept far more tokens than math or code, whose token distributions diverge from the training mix. There was no way to see a draft's acceptance profile across task types without manually invoking each script once per dataset and collating the numbers by hand. Add bench_sweep.py: drives either engine through several datasets in one pass and reports a per-dataset table plus a completed-weighted aggregate. - DatasetSpec + a built-in default suite: the four benchmarks the EAGLE / EAGLE-2 papers report acceptance/speedup on (MT-Bench, HumanEval, GSM8K, Alpaca). --datasets-config swaps in a custom YAML list; --datasets narrows to a subset. One dataset failing to load or benchmark is reported as an error row and does not abort the sweep, except a bad SHARED flag (e.g. --num-prompts 0), which is validated once up front and raises immediately instead of being reported as N duplicate per-dataset failures. - Flags an SGLang-specific caveat: avg_spec_accept_length is a server-cumulative running average with no reset/delta API, so sweeping more than one dataset against one live SGLang server contaminates every dataset after the first with prior traffic. A warning is logged when this applies; vLLM's counter-diffing approach has no such caveat. Supporting changes to make the sweep possible without duplicating logic: - bench_common._load_prompts gains a --prompt-column path (a raw text field wrapped into a single-turn user message) alongside the existing --messages-column path, since none of the four default datasets are chat-messages-shaped. Also available directly on bench_sglang/bench_vllm for a single custom dataset. bench_common._report_summary factors out the print-JSON-and-optionally-write-file tail bench_sglang/bench_vllm shared. - bench_sglang.py / bench_vllm.py: split _run into _run_summary (pure compute, returns the metrics dict or None) + _run (unchanged external behavior), so the sweep can drive one dataset at a time without the printing / --output-json side effects. Verified byte-for-byte behavior preserved against the existing test suites. - tools/lint_example_yamls.py: exempt the new dataset-list example config (not a recipe, launched via --datasets-config). Implements the multi-dataset benchmark suite item of NVIDIA-NeMo#2958. Signed-off-by: khazic <khazzz1c@gmail.com>
khazic
force-pushed
the
khazic/feat/spec-decode-bench-suite
branch
from
4049fae to
75e1b62
Compare
| # The classic EAGLE / EAGLE-2 paper suite: multi-turn chat (MT-Bench), code | ||
| # (HumanEval), math (GSM8K), and single-turn instruction-following (Alpaca). | ||
| DEFAULT_DATASET_PRESETS: tuple[DatasetSpec, ...] = ( | ||
| DatasetSpec(name="mt_bench", input_data="HuggingFaceH4/mt_bench_prompts", split="train", prompt_column="turns"), |
There was a problem hiding this comment.
Thank you for including MT-Bench in the default suite—this is a very useful addition. I noticed one schema mismatch that prevents this entry from running as shipped: the current Hugging Face dataset exposes the two-turn sequence under prompt, rather than turns. With this preset, _load_prompts returns zero prompts and the default sweep records MT-Bench as an error row. Could we update both this preset and the example YAML to use prompt_column="prompt", and add a representative schema test? Since _extract_prompt_text currently keeps only the first item, it would also be helpful either to describe this as a first-turn MT-Bench workload or to preserve the full two-turn flow so the multi-turn description remains accurate.
There was a problem hiding this comment.
Good catch, fixed. The preset and example YAML now use prompt_column: prompt (the real column on HuggingFaceH4/mt_bench_prompts); _extract_prompt_text reduces the two-turn list to its first turn, so this is a first-turn MT-Bench workload now, and I've updated the comments/docstring to say so. Added a schema test that the presets carry the right columns.
| max_new_tokens=512, | ||
| ), | ||
| DatasetSpec(name="gsm8k", input_data="openai/gsm8k", dataset_name="main", split="test", prompt_column="question"), | ||
| DatasetSpec(name="alpaca", input_data="tatsu-lab/alpaca", split="train", prompt_column="instruction"), |
There was a problem hiding this comment.
The Alpaca preset is a helpful inclusion. One concern is that selecting only instruction omits the input field, which carries required task context for many rows; the dataset card notes that roughly 40% of examples include it. In live validation, 23 of the first 64 rows had non-empty input—for example, “Identify the odd one out” loses “Twitter, Instagram, Telegram.” That leaves a substantial portion of the benchmark prompts incomplete and can make the acceptance statistics misleading. Could we add a small composition/template path that combines instruction and non-empty input without including the reference output, along with a regression test for both forms?
There was a problem hiding this comment.
Agreed. Added DatasetSpec.prompt_context_column: _load_prompts appends a non-empty secondary field (Alpaca's input) to the instruction, blank-line separated, and never reads the reference output. The alpaca preset and example YAML now set it. Added regression tests for both forms (input present -> joined, input empty/absent -> bare instruction).
| if summary is None: | ||
| results.append({"dataset": spec.name, "error": "no usable prompts loaded"}) | ||
| continue | ||
| results.append({"dataset": spec.name, **summary}) |
There was a problem hiding this comment.
The per-dataset isolation here is thoughtfully designed. There is one failure path where the sweep currently reports success: _run_workload intentionally catches request-level HTTP errors, so when every request fails, the engine returns a summary with completed=0 instead of raising or returning None. This row is then counted in num_datasets_ok, and a sweep where every request failed exits with status 0. Could we treat summaries with zero completed requests as error rows (or have the engine return None in that case) and add an all-requests-failed regression test? That would keep both the aggregate and process exit status aligned with the actual benchmark outcome.
There was a problem hiding this comment.
Fixed. _run_sweep now records a summary with completed < 1 as an error row ("all requests failed (0 completed)"), so it drops out of num_datasets_ok and the exit status reflects that nothing actually ran. Added an all-requests-failed regression test.
Address review on the multi-dataset bench sweep: - MT-Bench: HuggingFaceH4/mt_bench_prompts stores prompts under 'prompt' (a two-turn list), not 'turns'; the old preset loaded zero prompts. Use prompt_column='prompt' (first turn), in both the preset and example YAML. - Alpaca: ~40% of rows carry task context in a separate 'input' field that the instruction-only preset dropped. Add DatasetSpec.prompt_context_column, which _load_prompts appends to the prompt (blank-line separated) when non-empty; the reference 'output' is never read. Wire it for the alpaca preset / example YAML. - Sweep status: _run_workload swallows per-request HTTP errors, so a dataset whose every request failed returned a completed=0 summary that still counted in num_datasets_ok and let the sweep exit 0. Record completed<1 as an error row. Add unit tests for all three. Signed-off-by: khazic <khazzz1c@gmail.com>
| ``--messages-column`` expects, so each reads a raw text field instead | ||
| (``bench_common``'s ``--prompt-column`` path: the field is wrapped into a | ||
| fresh single-turn user message; a list value, e.g. MT-Bench's multi-turn | ||
| ``turns`` column, uses its first entry). |
There was a problem hiding this comment.
Thank you for addressing all three findings so quickly. I rechecked the updated presets and failure handling, and those fixes look good. I noticed just one small documentation follow-up: this module docstring still describes MT-Bench as “multi-turn chat” and refers to a “turns” column, while the current dataset uses the “prompt” column and this sweep intentionally benchmarks only the first turn. The README repeats the “multi-turn chat” wording. Would you mind updating both places to say “MT-Bench (first turn)” and replacing “turns” with “prompt”? Everything else looks good from my re-review—thank you again!
There was a problem hiding this comment.
Done, thanks. Updated both the module docstring and examples/speculative/README.md to say "MT-Bench (first turn)" and to describe the two-turn prompt column (dropping the stale "multi-turn chat" / "turns" wording). Pushed in d2908e0.
There was a problem hiding this comment.
Everything else looks good from your re-review
Follow-up to the review: the module docstring and examples/speculative/README.md still called MT-Bench 'multi-turn chat' and referenced the old 'turns' column. The sweep reads HuggingFaceH4/mt_bench_prompts' 'prompt' column and uses only the first turn, so say 'MT-Bench (first turn)' and 'two-turn prompt column'. Signed-off-by: khazic <khazzz1c@gmail.com>
3 participants
What does this PR do?
Implements the "multi-dataset acceptance benchmark suite" item of #2958.
Problem
bench_sglang.py/bench_vllm.pyeach measure acceptance length / speedup against ONE dataset per invocation. The same draft's acceptance rate varies sharply by task type (conversational data tends to accept far more tokens than math or code), so getting a draft's acceptance profile across task types meant manually invoking each script once per dataset and collating the numbers by hand.bench_sweep.pyDrives either engine (
--engine sglang|vllm) through several datasets in one pass and reports a per-dataset table plus a completed-weighted aggregate.--prompt-columnpath (see below).--datasets-config <path.yaml>swaps in a custom dataset list (seeexamples/speculative/bench_sweep/spec_bench_datasets.yaml, which mirrors the built-in default and doubles as a template);--datasets <name...>narrows to a subset.--num-prompts 0) is validated once up front instead, and raises immediately rather than being reported as N duplicate per-dataset failures.avg_spec_accept_lengthis a server-cumulative running average with no reset/delta API, so sweeping more than one dataset against one live SGLang server contaminates every dataset after the first with prior traffic. A warning is logged when this applies (--engine vllm's counter-diffing approach has no such caveat). Documented in the module docstring and README.Supporting changes (to avoid duplicating logic)
bench_common._load_promptsgains a--prompt-columnpath (a raw text field wrapped into a single-turn user message) alongside the existing--messages-columnpath, since none of the four default datasets are chat-messages-shaped. Also available directly onbench_sglang.py/bench_vllm.pyfor a single custom dataset, not just through the sweep.bench_common._report_summaryfactors out the print-JSON-and-optionally-write-file tail the two engine scripts shared.bench_sglang.py/bench_vllm.py: split_runinto_run_summary(pure compute, returns the metrics dict orNone) +_run(unchanged external behavior), so the sweep can drive one dataset at a time without the printing /--output-jsonside effects. Verified byte-for-byte behavior preserved against the existing test suites (all 57 pre-existing tests pass unchanged).tools/lint_example_yamls.py: exempt the new dataset-list example config (not a recipe, consumed via--datasets-config).Validation
tests/unit_tests/speculative/(94 new:test_bench_sweep.py+test_bench_common.py's new--prompt-columncoverage), plus the example-YAML linter suite.bench_sglang/bench_vllmtests use (_run_summarymonkeypatched per engine, no live server).bench_sglang/bench_vllmalready exercise (unchanged) in production.