nemo_automodel.components.speculative.streaming.producer

Streaming producer for speculative-decoding draft training.

The :class:FeatureProducer wraps a single target-backend forward pass and ships its output through a :class:FeatureStore. The trainer-side :class:~nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader re-materializes the same data as an :class:~nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch and yields it to the trainer loop.

The first producer wraps the existing :class:~nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel so the colocated path’s numerical output travels bit-for-bit into the streaming pipeline. Future producers (out-of-process SGLang, NCCL remote) plug in behind the same FeatureProducer API without touching the trainer.

Module Contents

Classes

Name	Description
`FeatureProducer`	Run a target backend once and put the supervision into a :class:`FeatureStore`.

Functions

Name	Description
`_resolve_algorithm`	Pick the producer-side :class:`FeatureAlgorithm` for a backend.

Data

__all__

logger

API

class nemo_automodel.components.speculative.streaming.producer.FeatureProducer(
    target_backend: nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend,
    store: nemo_automodel.components.speculative.streaming.store.FeatureStore,
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm | None = None,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    sample_id_factory: typing.Callable[[int], str] | None = None
)

Run a target backend once and put the supervision into a :class:FeatureStore.

Thread safety: produce is not reentrant with itself on the same producer instance; one thread should drive produce at a time. Multiple producers over the same store are fine; the store’s lock serializes the puts.

Parameters:

target_backend

Eagle3TargetBackend

an object exposing generate_batch(input_ids, attention_mask, loss_mask, ...) -> Eagle3TargetBatch plus the get_input_embeddings and set_vocab_mapping accessors from :class:~nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.

store

FeatureStore

The :class:FeatureStore every produce call writes into.

run_id

str

Stable across the whole training run; mirrored onto every :class:SampleRef.

algorithm

FeatureAlgorithm | NoneDefaults to None

Forced :class:FeatureAlgorithm; defaults to picking by the backend’s runtime type.

target_model_version

strDefaults to '0'

Monotonically increasing identifier of the target-model weights; surfaced on each :class:SampleRef.

draft_weight_version

strDefaults to '0'

Same idea for the draft model’s weights.

sample_id_factory

Callable[[int], str] | NoneDefaults to None

Callable that produces the per-call sample id; defaults to "sample-{run_id}-{n}". Wrap with a hash of the input batch if you want stable ids across runs.

_algorithm

_sample_id_factory

= sample_id_factory or self._default_sample_id

_sample_seq

= 0

nemo_automodel.components.speculative.streaming.producer.FeatureProducer._default_sample_id(
    n: int
) -> str

nemo_automodel.components.speculative.streaming.producer.FeatureProducer.close() -> None

Release any backend resources; the producer itself holds none.

nemo_automodel.components.speculative.streaming.producer.FeatureProducer.discard(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef
) -> None

Evict a produced sample whose ref will not be enqueued.

Used when a produced ref is dropped before it reaches the queue: a spec-mismatch abort in :meth:produce, or the async pipeline shutting down while a blocking put is in flight. get + release drops the last handle so the store frees the sample; without it the sample’s bytes (or its shared-dir file) stay resident with no consumer able to release them.

Parameters:

ref

SampleRef

The reference whose backing sample should be evicted.

nemo_automodel.components.speculative.streaming.producer.FeatureProducer.get_input_embeddings() -> torch.nn.Module

Expose the wrapped backend’s input-embedding module.

Mirrors :meth:Eagle3TargetBackend.get_input_embeddings so the draft can seed its input embeddings from the target before the first streamed batch arrives.

nemo_automodel.components.speculative.streaming.producer.FeatureProducer.produce(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef

Run the wrapped backend once and stash its supervision into the store.

Parameters:

input_ids

torch.Tensor

Tensor of shape [batch, sequence], torch.long.

attention_mask

torch.Tensor

Tensor of shape [batch, sequence], torch.long.

loss_mask

torch.Tensor

Tensor of shape [batch, sequence], torch.long.

position_ids

torch.Tensor | NoneDefaults to None

Per-document position ids when packing is on; None for non-packed. Carried through unchanged to the :class:Eagle3TargetBatch so the trainer’s loss path stays bit-identical to the colocated path.

seq_lens

torch.Tensor | NoneDefaults to None

Per-document sequence lengths when packing is on; None for non-packed.

doc_remaining

torch.Tensor | NoneDefaults to None

[batch, sequence] long tensor gating cross-document TTT supervision when packing is on; None otherwise.

Returns: SampleRef

class:SampleRef carrying a tensor-free description of

Raises:

MemoryError: forwarded from the store when the put would exceed max_samples or max_bytes. Producers can wait for the store to drain through the queue’s backpressure path (put_blocks_until_below).

nemo_automodel.components.speculative.streaming.producer.FeatureProducer.set_vocab_mapping(
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor
) -> None

Thread the draft-vocab mapping through to the wrapped backend.

The colocated backend ignores the mapping and projects the full-vocab logits trainer-side; a future draft-vocab encoding producer supplies the mapping here so the target can precompute target_probs + position_mask and the wire never carries a full-vocab tensor.

nemo_automodel.components.speculative.streaming.producer._resolve_algorithm(
    backend: object
) -> nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm

Pick the producer-side :class:FeatureAlgorithm for a backend.

Dispatch on the backend’s base class, not its class-name string: every Eagle3 target variant (HF, runner, vLLM, SGLang, remote) subclasses :class:Eagle3TargetBackend, so this survives renames, subclassing, and wrapping. Adding a new algorithm is a one-line isinstance branch here.

nemo_automodel.components.speculative.streaming.producer.__all__ = ['FeatureProducer']

nemo_automodel.components.speculative.streaming.producer.logger = logging.getLogger(__name__)