> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-speculative-eagle-msd-MSDStepMetrics)                       | Aggregated metrics from one multimodal speculative-decoding training step. |
| [`MSDTrainerModule`](#nemo_automodel-components-speculative-eagle-msd-MSDTrainerModule)                   | Draft-side training module for multimodal speculative decoding.            |
| [`MultimodalEagleDraftModel`](#nemo_automodel-components-speculative-eagle-msd-MultimodalEagleDraftModel) | EAGLE-1/2 draft model with modality-aware input feature construction.      |

### API

```python
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.

```python
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.

```python
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:**

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

**Returns:** `torch.Tensor`

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

```python
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:**

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

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

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

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

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

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

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

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

**Bases:** [LlamaEagleDraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel)

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

```python
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.

```python
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:**

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.

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

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

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

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

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