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.

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
)

Dataclass

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

accuracy

Tensor

hidden_loss

Tensor

loss

Tensor

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: float = 0.0
)

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.