nemo_automodel.components.speculative.bench_sweep

Sweep a trained drafter’s acceptance-length benchmark across multiple datasets.

bench_sglang / bench_vllm measure ONE workload per invocation. The same draft can behave very differently across task types — conversational data (ShareGPT / MT-Bench) tends to have far higher acceptance than math (GSM8K) or code (HumanEval), whose token distributions diverge sharply from the training mix. This script drives the same server through several named datasets in one pass and reports a per-dataset table plus a completed-weighted aggregate, instead of the user invoking bench_sglang/bench_vllm once per dataset and collating the output by hand.

Default dataset suite — the four benchmarks the EAGLE / EAGLE-2 papers report acceptance / speedup numbers on: MT-Bench (first turn), HumanEval (code), GSM8K (math), and Alpaca (single-turn instruction-following). None of these ship a chat-messages column the way bench_sglang/bench_vllm’s --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 two-turn prompt column, uses its first entry).

Override the default suite with --datasets-config <path.yaml>: a YAML list of entries, each {name, input_data, split, dataset_name, messages_column | prompt_column, max_new_tokens} (dataset_name/max_new_tokens optional; exactly one of messages_column/prompt_column required). --datasets further narrows either suite to a name subset.

Typical usage (after serve_sglang launches the drafter on port 30000)::

python -m nemo_automodel.components.speculative.bench_sweep
—engine sglang —server http://localhost:30000
—model meta-llama/Llama-3.1-8B-Instruct

Add --baseline-server for the speedup column, --engine vllm for a vLLM server, and --datasets-config to point at a custom dataset list. One dataset failing to load or benchmark (bad HF id, unreachable server for that request, …) does not abort the sweep — it is reported as an error row and excluded from the aggregate.

CAVEAT (--engine sglang with more than one dataset): SGLang’s avg_spec_accept_length is a server-cumulative running average with no reset/delta API (see bench_sglang’s module docstring), so sweeping N>1 datasets against ONE live SGLang server means every dataset after the first reports a blend with prior datasets’ traffic, not an independent number. A warning is logged when this applies. Restart the server between datasets for independent numbers, or use --engine vllm, which snapshots and diffs its Prometheus counters per dataset and has no such caveat.

Module Contents

Classes

Name	Description
`DatasetSpec`	One dataset to sweep: how to load it and where its prompt text lives.

Functions

Name	Description
`_aggregate`	Completed-weighted rollup of accept_length / acceptance_rate / throughput / speedup.
`_build_parser`	-
`_dataset_args`	Clone the shared server/model/workload args, overridden with one dataset’s spec.
`_fmt`	-
`_load_dataset_specs`	Return the sweep’s dataset list: the built-in suite, or a `--datasets-config` YAML override.
`_print_table`	Hand-rolled fixed-width table: one row per dataset, an error row for failures, then the aggregate.
`_row_line`	Format one fixed-width row (a dataset result or the aggregate) sharing one column layout.
`_run`	Async driver: sweep every selected dataset, print the table, optionally write JSON.
`_run_sweep`	Run one benchmark per dataset spec; one dataset’s failure does not abort the sweep.
`_select_datasets`	Narrow `specs` to `--datasets`, preserving the sweep’s declared order.
`_warn_if_sglang_stats_will_be_cumulative`	SGLang’s `avg_spec_accept_length` is a server-cumulative running average (see bench_sglang.py’s
`_weighted_mean`	Completed-weighted mean of `key` over already-filtered (error-free) results.
`main`	CLI entry point. Parses `argv` and returns the process exit code.

Data

DEFAULT_DATASET_PRESETS

_ENGINE_MODULES

_ENGINE_VALIDATORS

logger

API

class nemo_automodel.components.speculative.bench_sweep.DatasetSpec(
    name: str,
    input_data: str,
    split: str = 'train',
    dataset_name: str | None = None,
    messages_column: str | None = None,
    prompt_column: str | None = None,
    prompt_context_column: str | None = None,
    max_new_tokens: int | None = None
)

Dataclass

One dataset to sweep: how to load it and where its prompt text lives.

Exactly one of messages_column (an existing OpenAI-messages list) or prompt_column (a raw text field, wrapped into a single-turn user message) must be set — see bench_common._load_prompts.

prompt_context_column is an optional second raw-text field appended to prompt_column (separated by a blank line) when it is non-empty, for datasets whose task context lives in a separate column (e.g. Alpaca’s input). It is only valid alongside prompt_column.

dataset_name

str | None = None

input_data

str

max_new_tokens

int | None = None

messages_column

str | None = None

name

str

prompt_column

str | None = None

prompt_context_column

str | None = None

split

str = 'train'

nemo_automodel.components.speculative.bench_sweep.DatasetSpec.__post_init__()

nemo_automodel.components.speculative.bench_sweep._aggregate(
    results: list[dict[str, typing.Any]]
) -> dict[str, typing.Any]

Completed-weighted rollup of accept_length / acceptance_rate / throughput / speedup.

nemo_automodel.components.speculative.bench_sweep._build_parser() -> argparse.ArgumentParser

nemo_automodel.components.speculative.bench_sweep._dataset_args(
    base_args: argparse.Namespace,
    spec: nemo_automodel.components.speculative.bench_sweep.DatasetSpec
) -> argparse.Namespace

Clone the shared server/model/workload args, overridden with one dataset’s spec.

nemo_automodel.components.speculative.bench_sweep._fmt(
    value: typing.Any,
    digits: int = 3
) -> str

nemo_automodel.components.speculative.bench_sweep._load_dataset_specs(
    config_path: str | None
) -> list[nemo_automodel.components.speculative.bench_sweep.DatasetSpec]

Return the sweep’s dataset list: the built-in suite, or a --datasets-config YAML override.

The config’s top level must be a mapping with a datasets: list (not a bare top-level list) so the file passes the repo’s example-YAML linter, which requires every examples/ YAML to parse to a mapping.

nemo_automodel.components.speculative.bench_sweep._print_table(
    results: list[dict[str, typing.Any]],
    aggregate: dict[str, typing.Any]
) -> None

Hand-rolled fixed-width table: one row per dataset, an error row for failures, then the aggregate.

nemo_automodel.components.speculative.bench_sweep._row_line(
    name: str,
    row: dict[str, typing.Any]
) -> str

Format one fixed-width row (a dataset result or the aggregate) sharing one column layout.

nemo_automodel.components.speculative.bench_sweep._run(
    args: argparse.Namespace
) -> int

async

Async driver: sweep every selected dataset, print the table, optionally write JSON.

nemo_automodel.components.speculative.bench_sweep._run_sweep(
    args: argparse.Namespace,
    specs: list[nemo_automodel.components.speculative.bench_sweep.DatasetSpec]
) -> list[dict[str, typing.Any]]

async

Run one benchmark per dataset spec; one dataset’s failure does not abort the sweep.

nemo_automodel.components.speculative.bench_sweep._select_datasets(
    specs: list[nemo_automodel.components.speculative.bench_sweep.DatasetSpec],
    names: list[str] | None
) -> list[nemo_automodel.components.speculative.bench_sweep.DatasetSpec]

Narrow specs to --datasets, preserving the sweep’s declared order.

nemo_automodel.components.speculative.bench_sweep._warn_if_sglang_stats_will_be_cumulative(
    args: argparse.Namespace,
    num_datasets: int
) -> None

SGLang’s avg_spec_accept_length is a server-cumulative running average (see bench_sglang.py’s module docstring): it has no reset/delta API, so sweeping N>1 datasets against ONE live SGLang server means every dataset after the first reports a blend with prior datasets’ traffic, not an independent number. vLLM’s Prometheus counters are snapshotted before/after each dataset instead (see bench_vllm._run_summary), so this caveat is sglang-only.

nemo_automodel.components.speculative.bench_sweep._weighted_mean(
    ok_results: list[dict[str, typing.Any]],
    key: str
) -> float | None

Completed-weighted mean of key over already-filtered (error-free) results.

nemo_automodel.components.speculative.bench_sweep.main(
    argv: list[str] | None = None
) -> int

CLI entry point. Parses argv and returns the process exit code.

nemo_automodel.components.speculative.bench_sweep.DEFAULT_DATASET_PRESETS: tuple[DatasetSpec, ...] = (DatasetSpec(name='mt_bench', input_data='HuggingFaceH4/mt_bench_prompts', split...

nemo_automodel.components.speculative.bench_sweep._ENGINE_MODULES = {'sglang': bench_sglang, 'vllm': bench_vllm}

nemo_automodel.components.speculative.bench_sweep._ENGINE_VALIDATORS = {'sglang': bench_sglang._validate_args, 'vllm': bench_vllm._validate_workload_ar...

nemo_automodel.components.speculative.bench_sweep.logger = logging.getLogger(__name__)