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.

Module Contents

Classes

Name	Description
`VispecStepMetrics`	Aggregated metrics from one ViSpec training step.
`VispecTrainerModule`	Draft-side trainer for ViSpec stage-2 (vision-aware) training.

API

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

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

floatDefaults to 10.0

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

rank_loss_weight

floatDefaults to 0.1

Weight on the ListMLE term (ViSpec: 0.1).

rank_loss_topk

intDefaults to 10

Number of target tokens the ranking term covers.

mtp_steps

intDefaults to 1

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

feature_noise_config

FeatureNoiseConfig | NoneDefaults to 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.

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.

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

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/