nemo_automodel.components.speculative.regen_loop

On-policy data regeneration loop for EAGLE-3 draft training (train-with-decode).

EAGLE-3 teacher-forces the draft on a static dataset whose assistant turns were often written by a different model than the target being distilled, so the draft trains against off-distribution supervision (this is why the recipes ship an offline regenerate.py pass). This module runs that regeneration online, interleaved with training:

On a cadence (regen.every_steps), rank 0 launches a detached worker subprocess pinned to a reserved GPU (regen.cuda_visible_devices). The worker boots a plain vLLM OpenAI server for the frozen target and drives the existing regenerate machinery over a prompt slice, writing a fresh directory of parquet shards. Training never blocks; at most one cycle is in flight.
At the next epoch boundary the recipe checks for a completed cycle, and if one is ready all data-parallel ranks rebuild the training dataloader against the new shard directory in lockstep (the decision is broadcast from rank 0 so the ranks never diverge).

Because the EAGLE target is frozen, the regenerated distribution is stationary: the value here is pipelining (regenerate just-in-time instead of one giant upfront pass) and freshness (sampling new responses each cycle with temperature > 0), plus the plumbing a future draft-in-the-loop mode would reuse. The trainer process never imports vllm; the worker inherits the training env minus the torchrun/elastic variables (see decode_eval._worker_env).

Module Contents

Classes

Name	Description
`RegenConfig`	Resolved settings for the online on-policy regeneration loop.
`RegenRunner`	Rank-0 trainer-side orchestrator: launch regeneration at cadence, hand ready cycles to the recipe.

Functions

Name	Description
`_build_parser`	-
`_regenerate_argv`	Build the regenerate CLI args to run against the booted target server.
`_target_server_argv`	Build the plain (non-speculative) vLLM OpenAI server argv for the frozen target.
`main`	Worker entry: boot a target server, regenerate a shard dir, write the READY marker.
`resolve_regen_config`	Read the optional `regen:` block from `recipe_args`.

Data

_CONFIG_FILENAME

_DONE_FILENAME

_SHARDS_DIRNAME

_WORKER_LOG_FILENAME

logger

API

class nemo_automodel.components.speculative.regen_loop.RegenConfig(
    every_steps: int,
    cuda_visible_devices: str,
    target_model: str,
    input_data: str,
    output_dir: str,
    served_model_name: str = 'target',
    server_python: str | None = None,
    split: str = 'train',
    messages_column: str = 'messages',
    dataset_name: str | None = None,
    shard_size: int = 1000,
    concurrency: int = 32,
    max_new_tokens: int = 1024,
    temperature: float = 1.0,
    top_p: float = 1.0,
    shuffle_seed: int = 42,
    reasoning: str = 'none',
    port: int = 0,
    gpu_memory_utilization: float = 0.85,
    max_model_len: int | None = None,
    trust_remote_code: bool = False,
    timeout_s: float = 3600.0
)

Dataclass

Resolved settings for the online on-policy regeneration loop.

every_steps is the launch cadence in optimizer steps; a completed cycle is only swapped in at an epoch boundary (mid-epoch dataloader rebuilds would desync the sampler and the scheduler’s step count). cuda_visible_devices is the GPU (or comma list) reserved for the regeneration engine; training must not place work there. server_python is the interpreter used to launch the vLLM server (the training env need not have vllm installed); the regeneration client runs in the training env over HTTP.

concurrency

int = 32

cuda_visible_devices

str

dataset_name

str | None = None

every_steps

int

gpu_memory_utilization

float = 0.85

input_data

str

max_model_len

int | None = None

max_new_tokens

int = 1024

messages_column

str = 'messages'

output_dir

str

port

int = 0

reasoning

str = 'none'

served_model_name

str = 'target'

server_python

str | None = None

shard_size

int = 1000

shuffle_seed

int = 42

split

str = 'train'

target_model

str

temperature

float = 1.0

timeout_s

float = 3600.0

top_p

float = 1.0

trust_remote_code

bool = False

class nemo_automodel.components.speculative.regen_loop.RegenRunner(
    config: nemo_automodel.components.speculative.regen_loop.RegenConfig
)

Rank-0 trainer-side orchestrator: launch regeneration at cadence, hand ready cycles to the recipe.

At most one regeneration subprocess is alive at a time; a launch that lands while the previous cycle is still running is skipped (the next cadence boundary picks it up). A cycle is “ready” once its worker has written the READY marker; :meth:take_ready_shards returns the newest ready cycle’s shard directory exactly once, so the recipe can swap it into the dataloader.

_active_cycle

str | None = None

_consumed

set[str] = set()

_last_bucket

= 0

_launched_for_step

= -1

_proc

Popen | None = None

_worker_log_path

str | None = None

nemo_automodel.components.speculative.regen_loop.RegenRunner._cycle_dirs() -> list[tuple[int, str]]

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

nemo_automodel.components.speculative.regen_loop.RegenRunner._reap_finished_worker() -> None

Report a completed worker and release its Popen state.

nemo_automodel.components.speculative.regen_loop.RegenRunner._remove_cycle(
    cycle_path: str
) -> None

staticmethod

nemo_automodel.components.speculative.regen_loop.RegenRunner.due(
    global_step: int
) -> bool

Whether a new regeneration cycle should launch at this optimizer step.

nemo_automodel.components.speculative.regen_loop.RegenRunner.maybe_launch(
    global_step: int
) -> bool

Launch a regeneration worker if the cadence is due and none is running.

nemo_automodel.components.speculative.regen_loop.RegenRunner.resume_from_step(
    global_step: int
) -> None

Align the launch cadence to a restored global_step after a resume.

Without this a resumed run starts with _last_bucket == 0 and :meth:due fires immediately, relaunching a redundant cycle for a cadence region that already ran before the checkpoint. Superseded on-disk cycles left by the pre-crash run are reclaimed by :meth:take_ready_shards on the next swap (it frees every ready cycle but the newest), so no cycle path needs to be persisted across the checkpoint.

nemo_automodel.components.speculative.regen_loop.RegenRunner.shutdown() -> None

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

nemo_automodel.components.speculative.regen_loop.RegenRunner.take_ready_shards() -> str | None

Return the newest ready, not-yet-consumed cycle’s shard directory, or None.

A cycle is ready once its READY marker exists. Older ready cycles are marked consumed and skipped so the recipe always trains on the freshest regenerated data rather than replaying stale cycles.

nemo_automodel.components.speculative.regen_loop._build_parser()

nemo_automodel.components.speculative.regen_loop._regenerate_argv(
    cfg: nemo_automodel.components.speculative.regen_loop.RegenConfig,
    server: str,
    shards_dir: str
) -> list[str]

Build the regenerate CLI args to run against the booted target server.

nemo_automodel.components.speculative.regen_loop._target_server_argv(
    cfg: nemo_automodel.components.speculative.regen_loop.RegenConfig,
    port: int
) -> list[str]

Build the plain (non-speculative) vLLM OpenAI server argv for the frozen target.

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

Worker entry: boot a target server, regenerate a shard dir, write the READY marker.

nemo_automodel.components.speculative.regen_loop.resolve_regen_config(
    recipe_cfg: typing.Any,
    default_target: str,
    default_input_data: str | None,
    output_dir: str
) -> nemo_automodel.components.speculative.regen_loop.RegenConfig | None

Read the optional regen: 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 regenerating on a training GPU. Field defaults live on :class:RegenConfig only; the block just overrides the fields it sets.

nemo_automodel.components.speculative.regen_loop._CONFIG_FILENAME = 'regen_config.json'

nemo_automodel.components.speculative.regen_loop._DONE_FILENAME = 'READY'

nemo_automodel.components.speculative.regen_loop._SHARDS_DIRNAME = 'shards'

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

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