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`	Resolved settings for the periodic decode eval.
`DecodeEvalRunner`	Rank-0 trainer-side orchestrator: launch at cadence, collect results.

Functions

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

Data

_CONFIG_FILENAME

_RESULT_FILENAME

_SCRUBBED_ENV_KEYS

_SCRUBBED_ENV_PREFIXES

_WORKER_LOG_FILENAME

logger

API

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.

concurrency

int = 8

cuda_visible_devices

str

dataset_name

str | None = None

every_steps

int

gpu_memory_utilization

float = 0.8

input_data

str

max_model_len

int | None = None

max_new_tokens

int = 256

messages_column

str = 'messages'

num_prompts

int = 32

num_speculative_tokens

int

output_dir

str

port

int = 0

prompt_column

str | None = None

shuffle_seed

int = 42

split

str = 'train'

target_model

str

temperature

float = 0.0

timeout_s

float = 1800.0

top_p

float = 1.0

trust_remote_code

bool = False

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.

_collected

set[str] = set()

_last_bucket

= 0

_launched_for_step

= -1

_proc

Popen | None = None

_worker_log_path

str | None = None

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

staticmethod

Remove heavyweight weights after an eval no longer needs them.

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

Report a completed worker and release its Popen state.

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

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

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

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

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

Whether a new eval should launch at this optimizer step.

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.

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

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

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.

nemo_automodel.components.speculative.decode_eval._build_parser()

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.

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.

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.

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.

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:

draft_model

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.

out_dir

str

snapshot root; created if missing.

Returns: str

The model/consolidated directory path containing the export.

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

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

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.

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

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

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

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

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

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