> 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.vispec_core

ViSpec stage-2 draft-training objective (multi-token rollout + distribution losses).

Two things separate this from the EAGLE-1/2 objective in `core_v12.py`:

* **Self-rollout supervision.** After the first draft pass, the draft's *own*
  predicted hidden states are shifted right by one and fed back as its input
  features, `mtp_steps` times. Every rollout is supervised against the same
  target distribution. Training only on the target's fresh hidden states lets
  the draft lean on information it will not have at drafting depth > 1; rolling
  its own output back in removes that shortcut.
* **Distribution losses instead of hidden-state regression.** ViSpec drops the
  SmoothL1 hidden-state term entirely and supervises the *distribution*: an L1
  distance between the draft's and the target's full-vocab probabilities, plus a
  ListMLE ranking term over the target's top-k tokens.

Reference implementation: `vispec/train/main_mtp.py` in
[https://github.com/KangJialiang/ViSpec](https://github.com/KangJialiang/ViSpec).

## Module Contents

### Classes

| Name                                                                                                  | Description                                                    |
| ----------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`VispecStepMetrics`](#nemo_automodel-components-speculative-eagle-vispec_core-VispecStepMetrics)     | Aggregated metrics from one ViSpec training step.              |
| [`VispecTrainerModule`](#nemo_automodel-components-speculative-eagle-vispec_core-VispecTrainerModule) | Draft-side trainer for ViSpec stage-2 (vision-aware) training. |

### API

```python
class nemo_automodel.components.speculative.eagle.vispec_core.VispecStepMetrics(
    loss: torch.Tensor,
    prob_loss: torch.Tensor,
    rank_loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor
)
```

Dataclass

Aggregated metrics from one ViSpec training step.

**`accuracy`** `Tensor`

---

**`loss`** `Tensor`

---

**`prob_loss`** `Tensor`

---

**`rank_loss`** `Tensor`

---

**`valid_tokens`** `Tensor`

---

```python
class nemo_automodel.components.speculative.eagle.vispec_core.VispecTrainerModule(
    draft_model: torch.nn.Module,
    target_lm_head: torch.nn.Module,
    prob_loss_weight: float = 10.0,
    rank_loss_weight: float = 0.1,
    rank_loss_topk: int = 10,
    mtp_steps: int = 1,
    feature_noise_config: nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig | None = None
)
```

**Bases:** `Module`

Draft-side trainer for ViSpec stage-2 (vision-aware) training.

**Parameters:**

**`draft_model`** `nn.Module`

The ViSpec draft model being trained.

---

**`target_lm_head`** `nn.Module`

The frozen target `lm_head`, used to turn predicted
hidden states into logits. Held off the module registry so it is
not duplicated into the draft's state dict / DDP buckets.

---

**`prob_loss_weight`** `float` — default: 10.0

Weight on the probability-L1 term (ViSpec: 10.0).

---

**`rank_loss_weight`** `float` — default: 0.1

Weight on the ListMLE term (ViSpec: 0.1).

---

**`rank_loss_topk`** `int` — default: 10

Number of target tokens the ranking term covers.

---

**`mtp_steps`** `int` — default: 1

Number of self-rollout passes after the first draft pass.

---

**`feature_noise_config`** `FeatureNoiseConfig | None` — default: None

Train-only feature augmentation, applied to the
target features entering the *first* pass. ViSpec's stage 2 enables
the same sequence-scaled draw as stage 1; `None` disables it.

---

```python
nemo_automodel.components.speculative.eagle.vispec_core.VispecTrainerModule._lm_head_weight() -> torch.Tensor
```

Return the frozen target lm\_head weight as a plain local tensor.

An FSDP2-sharded target exposes a DTensor weight; the draft runs under
DDP with plain tensors, so it has to be gathered before `F.linear`.
The target is frozen, so the gathered result is cached: resolving it
inside the rollout loop would repeat a full `[vocab, hidden]`
all-gather once per rollout, per micro-batch.

```python
nemo_automodel.components.speculative.eagle.vispec_core.VispecTrainerModule.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Project hidden states through the frozen target `lm_head`.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., vocab].

```python
nemo_automodel.components.speculative.eagle.vispec_core.VispecTrainerModule.forward(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    image_mask: torch.Tensor
) -> nemo_automodel.components.speculative.eagle.vispec_core.VispecStepMetrics
```

Run one ViSpec training step (first pass plus `mtp_steps` self-rollouts).

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Tensor of shape \[1, sequence, hidden] -- target
embedding-layer output, shifted left by one position.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape \[1, sequence]; 1 for real tokens.

---

**`loss_mask`** `torch.Tensor`

Tensor of shape \[1, sequence]; 1 at supervised positions.

---

**`input_hidden_states`** `torch.Tensor`

Tensor of shape \[1, sequence, hidden] -- the
target's last hidden state, not shifted.

---

**`target_logits`** `torch.Tensor`

Tensor of shape \[1, sequence, vocab] -- the target's
logits, shifted left by one position.

---

**`image_mask`** `torch.Tensor`

Bool tensor of shape \[1, sequence] aligned with
`inputs_embeds`.

---

**Returns:** `VispecStepMetrics`

VispecStepMetrics with scalar `loss`/`prob_loss`/`rank_loss`/