nemo_automodel.components.speculative.eagle.msd

Multimodal speculative decoding draft model and training objective.

MSD extends EAGLE-1/2 feature drafting to VLM targets. Text positions use the usual concatenation of target features and next-token embeddings. Image positions instead pass the target VLM’s already-computed image embeddings directly into the draft transformer, preserving the non-causal relationship between visual patches.

Module Contents

Classes

Name	Description
`MSDStepMetrics`	Aggregated metrics from one multimodal speculative-decoding training step.
`MSDTrainerModule`	Draft-side training module for multimodal speculative decoding.
`MultimodalEagleDraftModel`	EAGLE-1/2 draft model with modality-aware input feature construction.

API

class nemo_automodel.components.speculative.eagle.msd.MSDStepMetrics(
    loss: torch.Tensor,
    hidden_loss: torch.Tensor,
    token_loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor
)

Dataclass

Aggregated metrics from one multimodal speculative-decoding training step.

accuracy

Tensor

hidden_loss

Tensor

loss

Tensor

token_loss

Tensor

valid_tokens

Tensor

class nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    hidden_loss_weight: float = 1.0,
    token_loss_weight: float = 0.1
)

Bases: Module

Draft-side training module for multimodal speculative decoding.

hidden_loss_fn

= nn.SmoothL1Loss(reduction='none')

nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

Project predicted hidden states through the frozen target language head.

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

Returns: torch.Tensor

Tensor of shape […, vocab], containing target-vocabulary logits.

nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.forward(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    image_mask: torch.Tensor
) -> nemo_automodel.components.speculative.eagle.msd.MSDStepMetrics

Run one multimodal speculative-decoding training step.

Parameters:

inputs_embeds

torch.Tensor

Tensor of shape [batch, sequence, hidden], the target VLM input embeddings shifted left by one position.

attention_mask

torch.Tensor

Tensor of shape [batch, sequence], where one denotes a real token and zero denotes padding.

loss_mask

torch.Tensor

Bool tensor of shape [batch, sequence], true where the shifted next-token target is supervised.

input_hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden], the target’s unshifted final hidden states used as draft features.

target_hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden], the target’s final hidden states shifted left by one position.

target_logits

torch.Tensor

Tensor of shape [batch, sequence, vocab], the frozen target logits shifted left by one position.

image_mask

torch.Tensor

Bool tensor of shape [batch, sequence], true where inputs_embeds contains a projected image token.

Returns: MSDStepMetrics

MSDStepMetrics containing scalar loss terms, first-token accuracy, and

class nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel()

Bases: LlamaEagleDraftModel

EAGLE-1/2 draft model with modality-aware input feature construction.

nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.copy_embeddings_from_target(
    target_embeddings: torch.nn.Module
) -> None

Raise: MSD drafts consume target embeddings instead of holding their own.

nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.forward(
    inputs_embeds: torch.Tensor,
    target_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    image_mask: torch.Tensor
) -> torch.Tensor

Predict next-position target hidden states from multimodal features.

Parameters:

inputs_embeds

torch.Tensor

Tensor of shape [batch, sequence, hidden], containing the target VLM input embeddings shifted one position to the left. Image positions already contain projected vision features.

target_hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden], the frozen target language model’s unshifted final hidden states.

attention_mask

torch.Tensor

Tensor of shape [batch, sequence], where one denotes a real token and zero denotes padding.

image_mask

torch.Tensor

Bool tensor of shape [batch, sequence], aligned with inputs_embeds and true for projected image-token positions.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden], the draft predictions for

nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.freeze_embeddings() -> None

Raise: MSD drafts have no token embedding table to freeze.