> 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.decode_eval

Periodic real-acceptance-length eval during draft training (decode eval).

Training-time draft metrics (loss, top-1 accuracy, the simulated `tau_sim`)
are proxies: the acceptance length that matters is produced by the draft and
the target interacting inside a real speculative-decoding engine. This module
closes that gap without touching the training step:

* On a cadence (`decode_eval.every_steps`), rank 0 snapshots the current
  draft weights to disk and launches a detached **worker subprocess** pinned to
  a reserved GPU (`decode_eval.cuda_visible_devices`). Training continues
  immediately; at most one eval is in flight.
* The worker (this module's CLI) packages the snapshot through the existing
  `serve_vllm` conversion, boots a vLLM OpenAI server as its child, drives a
  fixed prompt set through it with the existing `bench_vllm` workload runner,
  and writes a JSON result (real `accept_length` from the engine's
  spec-decode counters) next to the snapshot.
* The trainer's logging block collects finished results and logs them as
  `train/tau_real` (with `train/tau_real_step` marking the optimizer step
  the evaluated snapshot was taken at).

The trainer process never imports vllm; the worker inherits the training env
minus the torchrun/elastic variables, so the engine initializes as a plain
single-process job on the reserved GPU.

## Module Contents

### Classes

| Name                                                                                      | Description                                                           |
| ----------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`DecodeEvalConfig`](#nemo_automodel-components-speculative-decode_eval-DecodeEvalConfig) | Resolved settings for the periodic decode eval.                       |
| [`DecodeEvalRunner`](#nemo_automodel-components-speculative-decode_eval-DecodeEvalRunner) | Rank-0 trainer-side orchestrator: launch at cadence, collect results. |

### Functions

| Name                                                                                                          | Description                                                                           |
| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_bench_args`](#nemo_automodel-components-speculative-decode_eval-_bench_args)                               | Assemble the bench\_vllm.\_run\_summary argument namespace.                           |
| [`_build_parser`](#nemo_automodel-components-speculative-decode_eval-_build_parser)                           | -                                                                                     |
| [`_resolve_worker_port`](#nemo_automodel-components-speculative-decode_eval-_resolve_worker_port)             | Choose an unused loopback port, or fail if an explicit port is occupied.              |
| [`_serve_argv`](#nemo_automodel-components-speculative-decode_eval-_serve_argv)                               | Build the vLLM api\_server argv for the snapshot via the serve\_vllm library surface. |
| [`_wait_for_server`](#nemo_automodel-components-speculative-decode_eval-_wait_for_server)                     | Poll the vLLM server's /health until it responds or dies/times out.                   |
| [`_worker_env`](#nemo_automodel-components-speculative-decode_eval-_worker_env)                               | Training env minus torchrun/elastic state, pinned to the reserved GPU.                |
| [`export_draft_snapshot`](#nemo_automodel-components-speculative-decode_eval-export_draft_snapshot)           | Write the draft's current weights as a serve-ready consolidated export.               |
| [`main`](#nemo_automodel-components-speculative-decode_eval-main)                                             | Worker entry: boot the engine on the snapshot, bench it, write the result JSON.       |
| [`resolve_decode_eval_config`](#nemo_automodel-components-speculative-decode_eval-resolve_decode_eval_config) | Read the optional `decode_eval:` block from `recipe_args`.                            |

### Data

[`_CONFIG_FILENAME`](#nemo_automodel-components-speculative-decode_eval-_CONFIG_FILENAME)

[`_RESULT_FILENAME`](#nemo_automodel-components-speculative-decode_eval-_RESULT_FILENAME)

[`_SCRUBBED_ENV_KEYS`](#nemo_automodel-components-speculative-decode_eval-_SCRUBBED_ENV_KEYS)

[`_SCRUBBED_ENV_PREFIXES`](#nemo_automodel-components-speculative-decode_eval-_SCRUBBED_ENV_PREFIXES)

[`_WORKER_LOG_FILENAME`](#nemo_automodel-components-speculative-decode_eval-_WORKER_LOG_FILENAME)

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

### API

```python
class nemo_automodel.components.speculative.decode_eval.DecodeEvalConfig(
    every_steps: int,
    cuda_visible_devices: str,
    target_model: str,
    input_data: str,
    output_dir: str,
    num_speculative_tokens: int,
    num_prompts: int = 32,
    concurrency: int = 8,
    max_new_tokens: int = 256,
    temperature: float = 0.0,
    top_p: float = 1.0,
    messages_column: str = 'messages',
    prompt_column: str | None = None,
    split: str = 'train',
    dataset_name: str | None = None,
    shuffle_seed: int = 42,
    port: int = 0,
    gpu_memory_utilization: float = 0.8,
    max_model_len: int | None = None,
    trust_remote_code: bool = False,
    timeout_s: float = 1800.0
)
```

Dataclass

Resolved settings for the periodic decode eval.

`every_steps` is the launch cadence in optimizer steps; launches are
checked at the recipe's logging points, so it should be a multiple of
`log_every_steps` (a non-multiple fires at the first log point past the
boundary). `cuda_visible_devices` is the GPU (or comma list) reserved for
the eval engine; training must not place work there.

```python
class nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner(
    config: nemo_automodel.components.speculative.decode_eval.DecodeEvalConfig
)
```

Rank-0 trainer-side orchestrator: launch at cadence, collect results.

At most one eval subprocess is alive at a time; a launch that lands while
the previous eval is still running is skipped (the next cadence boundary
picks it up). Results are one-line JSON files the worker writes on success;
`collect()` returns the not-yet-seen ones in step order.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner._cleanup_snapshot(
    step_dir: str
) -> None
```

staticmethod

Remove heavyweight weights after an eval no longer needs them.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner._reap_finished_worker() -> None
```

Report a completed worker and release its Popen state.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner._step_dirs() -> list[tuple[int, str]]
```

Return valid step directories in numeric order, ignoring stray names.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner.collect() -> list[dict]
```

Return finished, not-yet-reported eval results in snapshot-step order.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner.due(
    global_step: int
) -> bool
```

Whether a new eval should launch at this optimizer step.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner.maybe_launch(
    global_step: int,
    draft_model
) -> bool
```

Snapshot the draft and launch the worker if the cadence is due and no eval is running.

```python
nemo_automodel.components.speculative.decode_eval.DecodeEvalRunner.shutdown() -> None
```

Terminate a still-running eval (its vLLM child dies with the session).

```python
nemo_automodel.components.speculative.decode_eval._bench_args(
    cfg: nemo_automodel.components.speculative.decode_eval.DecodeEvalConfig,
    server: str
) -> argparse.Namespace
```

Assemble the bench\_vllm.\_run\_summary argument namespace.

```python
nemo_automodel.components.speculative.decode_eval._build_parser()
```

```python
nemo_automodel.components.speculative.decode_eval._resolve_worker_port(
    configured_port: int
) -> int
```

Choose an unused loopback port, or fail if an explicit port is occupied.

```python
nemo_automodel.components.speculative.decode_eval._serve_argv(
    cfg: nemo_automodel.components.speculative.decode_eval.DecodeEvalConfig,
    draft: str
) -> list[str]
```

Build the vLLM api\_server argv for the snapshot via the serve\_vllm library surface.

```python
nemo_automodel.components.speculative.decode_eval._wait_for_server(
    server: str,
    proc: subprocess.Popen,
    timeout_s: float
) -> None
```

Poll the vLLM server's /health until it responds or dies/times out.

```python
nemo_automodel.components.speculative.decode_eval._worker_env(
    cuda_visible_devices: str
) -> dict[str, str]
```

Training env minus torchrun/elastic state, pinned to the reserved GPU.

```python
nemo_automodel.components.speculative.decode_eval.export_draft_snapshot(
    draft_model,
    out_dir: str
) -> str
```

Write the draft's current weights as a serve-ready consolidated export.

Produces `<out_dir>/model/consolidated/{model.safetensors, config.json}`,
the same layout the final checkpoint's consolidated export uses, so
`serve_vllm.resolve_draft_artifacts` can consume it directly (the
d2t/t2d vocab-mapping buffers ride along in the state dict).

Assumes the draft is replicated on the calling rank (the EAGLE-3 draft
trains under DDP, so rank 0 holds the full weights); a sharded draft would
snapshot an incomplete state dict.

**Parameters:**

the (unwrapped) draft `nn.Module`; its parameter and
buffer tensors are copied to CPU, so shapes/layouts are whatever the
draft's `state_dict()` reports.

snapshot root; created if missing.

**Returns:** `str`

The `model/consolidated` directory path containing the export.

```python
nemo_automodel.components.speculative.decode_eval.main(
    argv = None
) -> int
```

Worker entry: boot the engine on the snapshot, bench it, write the result JSON.

```python
nemo_automodel.components.speculative.decode_eval.resolve_decode_eval_config(
    recipe_cfg: typing.Any,
    default_target: str,
    default_input_data: str | None,
    default_num_speculative_tokens: int,
    output_dir: str
) -> nemo_automodel.components.speculative.decode_eval.DecodeEvalConfig | None
```

Read the optional `decode_eval:` block from `recipe_args`.

Returns `None` when the block is absent or `every_steps` is unset/0
(feature disabled). Raises on a partially-configured block so a typo'd GPU
reservation fails at setup rather than silently evaluating on a training
GPU. Field defaults live on :class:`DecodeEvalConfig` only; the block just
overrides the fields it sets.

```python
nemo_automodel.components.speculative.decode_eval._CONFIG_FILENAME = 'decode_eval_config.json'
```

```python
nemo_automodel.components.speculative.decode_eval._RESULT_FILENAME = 'result.json'
```

```python
nemo_automodel.components.speculative.decode_eval._SCRUBBED_ENV_KEYS = ('RANK', 'LOCAL_RANK', 'WORLD_SIZE', 'LOCAL_WORLD_SIZE', 'GROUP_RANK', 'GROUP_WO...
```

```python
nemo_automodel.components.speculative.decode_eval._SCRUBBED_ENV_PREFIXES = ('TORCHELASTIC_', 'PET_')
```

```python
nemo_automodel.components.speculative.decode_eval._WORKER_LOG_FILENAME = 'worker.log'
```

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