nemo_automodel.components.speculative.eagle.vispec_target

Target-model wrapper for ViSpec draft training on a vision-language target.

Where the EAGLE-1/2 wrapper hands the draft token ids, ViSpec hands it the target’s embedding-layer output: at image positions there is no token embedding to look up, only the vision tower’s projected features, and those features are exactly what the draft’s image adaptor compresses. The wrapper therefore returns the target’s layer-0 hidden states alongside the usual last-hidden-state / logits supervision, plus the image-token mask that tells the draft which positions to compress.

Module Contents

Classes

Name	Description
`HFVispecTargetModel`	Expose embedding-layer, last-hidden-state, and logit supervision from a VLM target.
`VispecTargetBatch`	Target-model outputs needed by :class:`VispecTrainerModule`.

API

class nemo_automodel.components.speculative.eagle.vispec_target.HFVispecTargetModel(
    model: torch.nn.Module,
    image_token_id: int
)

Expose embedding-layer, last-hidden-state, and logit supervision from a VLM target.

Parameters:

model

nn.Module

The frozen vision-language target model.

image_token_id

int

Token id the target uses as an image placeholder; every position holding it carries a vision feature instead of a token embedding.

_accepted_params

= frozenset(forward_params)

_extra_kwargs

image_token_id

= int(image_token_id)

model

= model.eval()

nemo_automodel.components.speculative.eagle.vispec_target.HFVispecTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    multimodal_inputs: torch.Tensor = {}
) -> nemo_automodel.components.speculative.eagle.vispec_target.VispecTargetBatch

Run the frozen target once and assemble the draft’s supervision.

Parameters:

input_ids

torch.Tensor

Tensor of shape [batch, sequence].

attention_mask

torch.Tensor

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

loss_mask

torch.Tensor

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

**multimodal_inputs

torch.TensorDefaults to {}

The processor’s vision tensors for this batch (e.g. pixel_values of shape [patches, patch_dim] and image_grid_thw of shape [images, 3] for Qwen2.5-VL). Keys the target’s forward does not declare are dropped.

Returns: VispecTargetBatch

VispecTargetBatch, with every tensor on the target’s device.

nemo_automodel.components.speculative.eagle.vispec_target.HFVispecTargetModel.get_input_embeddings() -> torch.nn.Module

Return the target model input embeddings.

nemo_automodel.components.speculative.eagle.vispec_target.HFVispecTargetModel.get_lm_head() -> torch.nn.Module

Return the target model lm_head.

class nemo_automodel.components.speculative.eagle.vispec_target.VispecTargetBatch(
    inputs_embeds: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    image_mask: torch.Tensor
)

Dataclass

Target-model outputs needed by :class:VispecTrainerModule.

attention_mask

Tensor

image_mask

Tensor

input_hidden_states

Tensor

inputs_embeds

Tensor

loss_mask

Tensor

target_logits

Tensor