> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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](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`](#nemo_automodel-components-speculative-bench_sweep-DatasetSpec) | One dataset to sweep: how to load it and where its prompt text lives. |

### Functions

| Name                                                                                                                                      | Description                                                                                        |
| ----------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- |
| [`_aggregate`](#nemo_automodel-components-speculative-bench_sweep-_aggregate)                                                             | Completed-weighted rollup of accept\_length / acceptance\_rate / throughput / speedup.             |
| [`_build_parser`](#nemo_automodel-components-speculative-bench_sweep-_build_parser)                                                       | -                                                                                                  |
| [`_dataset_args`](#nemo_automodel-components-speculative-bench_sweep-_dataset_args)                                                       | Clone the shared server/model/workload args, overridden with one dataset's spec.                   |
| [`_fmt`](#nemo_automodel-components-speculative-bench_sweep-_fmt)                                                                         | -                                                                                                  |
| [`_load_dataset_specs`](#nemo_automodel-components-speculative-bench_sweep-_load_dataset_specs)                                           | Return the sweep's dataset list: the built-in suite, or a `--datasets-config` YAML override.       |
| [`_print_table`](#nemo_automodel-components-speculative-bench_sweep-_print_table)                                                         | Hand-rolled fixed-width table: one row per dataset, an error row for failures, then the aggregate. |
| [`_row_line`](#nemo_automodel-components-speculative-bench_sweep-_row_line)                                                               | Format one fixed-width row (a dataset result or the aggregate) sharing one column layout.          |
| [`_run`](#nemo_automodel-components-speculative-bench_sweep-_run)                                                                         | Async driver: sweep every selected dataset, print the table, optionally write JSON.                |
| [`_run_sweep`](#nemo_automodel-components-speculative-bench_sweep-_run_sweep)                                                             | Run one benchmark per dataset spec; one dataset's failure does not abort the sweep.                |
| [`_select_datasets`](#nemo_automodel-components-speculative-bench_sweep-_select_datasets)                                                 | Narrow `specs` to `--datasets`, preserving the sweep's declared order.                             |
| [`_warn_if_sglang_stats_will_be_cumulative`](#nemo_automodel-components-speculative-bench_sweep-_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`](#nemo_automodel-components-speculative-bench_sweep-_weighted_mean)                                                     | Completed-weighted mean of `key` over already-filtered (error-free) results.                       |
| [`main`](#nemo_automodel-components-speculative-bench_sweep-main)                                                                         | CLI entry point. Parses `argv` and returns the process exit code.                                  |

### Data

[`DEFAULT_DATASET_PRESETS`](#nemo_automodel-components-speculative-bench_sweep-DEFAULT_DATASET_PRESETS)

[`_ENGINE_MODULES`](#nemo_automodel-components-speculative-bench_sweep-_ENGINE_MODULES)

[`_ENGINE_VALIDATORS`](#nemo_automodel-components-speculative-bench_sweep-_ENGINE_VALIDATORS)

[`logger`](#nemo_automodel-components-speculative-bench_sweep-logger)

### API

```python
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`.

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

```python
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.

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

```python
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.

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

```python
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.

```python
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.

```python
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.

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

async

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

```python
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.

```python
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.

```python
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.

```python
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.

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

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

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

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

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

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