nemo_automodel.components.speculative.eagle.msd_target

Frozen VLM target wrapper for multimodal speculative-decoding training.

Module Contents

Classes

Name	Description
`HFMSDTargetModel`	Expose image-aware feature supervision from a frozen Hugging Face VLM.
`MSDTargetBatch`	Frozen VLM supervision tensors consumed by :class:`MSDTrainerModule`.
`_InputEmbeddingCapture`	Capture the input embeddings supplied to a VLM language backbone.

API

class nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel(
    model: torch.nn.Module
)

Expose image-aware feature supervision from a frozen Hugging Face VLM.

image_token_id

= int(image_token_id)

language_backbone

model

= model.eval()

nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.generate_batch(
    loss_mask: torch.Tensor,
    model_inputs: dict[str, torch.Tensor]
) -> nemo_automodel.components.speculative.eagle.msd_target.MSDTargetBatch

Run a VLM and construct shifted multimodal draft supervision.

input_ids and attention_mask are read from model_inputs rather than taken separately, so the tensors the target actually runs on are the same ones the image mask and the returned attention mask are derived from. A second copy would misalign the draft’s modality routing without raising.

Parameters:

loss_mask

torch.Tensor

Bool tensor of shape [batch, sequence], already aligned to next-token logits. VLM collators produce this from labels != -100; they emit labels[:, 1:] against input_ids[:, :-1], so no further shift is applied here. Its final position is always dropped because the shifted supervision tensors zero-fill there.

model_inputs

dict[str, torch.Tensor]

Mapping of VLM forward inputs. It must include input_ids of shape [batch, sequence], with image-token placeholders at positions that receive projected vision features, and attention_mask of the same shape, where one denotes a real token and zero denotes padding. It also carries any vision tensors such as pixel_values and image_grid_thw.

Returns: MSDTargetBatch

MSDTargetBatch with shifted image-aware embeddings, target hidden

nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.get_lm_head() -> torch.nn.Module

Return the frozen VLM language-model head used by the draft loss.

class nemo_automodel.components.speculative.eagle.msd_target.MSDTargetBatch(
    inputs_embeds: torch.Tensor,
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    image_mask: torch.Tensor
)

Dataclass

Frozen VLM supervision tensors consumed by :class:MSDTrainerModule.

attention_mask

Tensor

image_mask

Tensor

input_hidden_states

Tensor

inputs_embeds

Tensor

loss_mask

Tensor

target_hidden_states

Tensor

target_logits

Tensor

class nemo_automodel.components.speculative.eagle.msd_target._InputEmbeddingCapture()

Capture the input embeddings supplied to a VLM language backbone.

value

Tensor | None = None

nemo_automodel.components.speculative.eagle.msd_target._InputEmbeddingCapture.__call__(
    _module: torch.nn.Module,
    _args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> None

Save the fused VLM embedding tensor passed to the language backbone.

Parameters:

_module

nn.Module

Language backbone receiving the VLM’s fused input embeddings.

_args

tuple[Any, ...]

Positional arguments accepted by the language backbone.

kwargs

dict[str, Any]

Keyword arguments accepted by the language backbone. Its inputs_embeds value is a Tensor of shape [batch, sequence, hidden] containing text and projected image embeddings.