nemo_automodel.components.speculative.eagle.core_v12

Core EAGLE-1 / EAGLE-2 draft-training logic.

Module Contents

Classes

Name	Description
`EagleStepMetrics`	Aggregated metrics from one EAGLE-1 / EAGLE-2 training step.
`EagleTrainerModule`	Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction.
`FeatureNoiseConfig`	Train-only uniform noise added to the target features fed to the draft.

API

class nemo_automodel.components.speculative.eagle.core_v12.EagleStepMetrics(
    loss: torch.Tensor,
    hidden_loss: torch.Tensor,
    token_loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor,
    rank_loss: torch.Tensor | None = None
)

Dataclass

Aggregated metrics from one EAGLE-1 / EAGLE-2 training step.

accuracy

Tensor

hidden_loss

Tensor

loss

Tensor

rank_loss

Tensor | None = None

token_loss

Tensor

valid_tokens

Tensor

class nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    hidden_loss_weight: float = 1.0,
    token_loss_weight: float = 0.1,
    feature_noise_config: nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig | None = None,
    rank_loss_weight: float = 0.0,
    rank_loss_topk: int = 10
)

Bases: Module

Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction.

hidden_loss_fn

= nn.SmoothL1Loss(reduction='none')

nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

Project predicted hidden states through the frozen target lm_head.

nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.core_v12.EagleStepMetrics

Run one EAGLE-1 / EAGLE-2 training step.

Per-token tensors are [B, T] (position_ids / doc_remaining included) except the [B, T, H] hidden states and [B, T, V] target_logits; seq_lens is [B, max_docs]. When packing is on, position_ids / seq_lens make the draft block-causal and per-document, and doc_remaining (real tokens after each slot within its document) gates supervision: a document’s last real token (doc_remaining == 0) is dropped because the wrapper’s global left-shift makes its target the next document’s first token.

class nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig(
    std: float = 0.2,
    reference_seq_len: int | None = 512
)

Dataclass

Train-only uniform noise added to the target features fed to the draft.

EAGLE’s reference implementation draws the perturbation as (rand_like(x) - 0.5) * std * reference_seq_len / T, where T is the effective (unpadded) sequence length of the batch. The half-width is therefore std / 2 * reference_seq_len / T: it is calibrated at reference_seq_len and shrinks as the sequence grows, so a longer context is not perturbed proportionally harder. ViSpec inherits this unchanged and enables it in both of its stages.

reference_seq_len=None drops the scaling and applies a fixed U(-std/2, std/2) at every length. That is the EAGLE paper’s wording (U(-0.1, 0.1), i.e. std=0.2) and what the EAGLE-1/2 recipe uses; it is 8x the reference half-width once T reaches 4096, which is why the ViSpec stages take the scaled form.

reference_seq_len

int | None = 512

std

float = 0.2

nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig.apply(
    features: torch.Tensor,
    attention_mask: torch.Tensor
) -> torch.Tensor

Return features perturbed by the uniform draw.

Parameters:

features

torch.Tensor

Tensor of shape [batch, sequence, hidden] containing target features handed to the draft.

attention_mask

torch.Tensor

Tensor of shape [batch, sequence]. Its longest unpadded row sets the scale, so padding does not shrink the perturbation.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden]. This is a new tensor;

nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig.fixed(
    half_width: float
) -> 'FeatureNoiseConfig'

classmethod

Build an unscaled U(-half_width, half_width) draw.

nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig.half_width(
    seq_len: int
) -> float

Return the symmetric noise half-width for a sequence of seq_len.