nemo_automodel.components.speculative.eagle.target

Target-model wrapper for minimal EAGLE-3 training.

Module Contents

Classes

Name	Description
`Eagle3TargetBatch`	Target-model supervision for one draft-training batch.
`HFEagle3TargetModel`	Co-located backend that captures three auxiliary hidden states from a causal LM.

Functions

Name	Description
`_shift_left_with_zero`	Shift a batched sequence tensor left and zero-fill the tail.
`_to_full_tensor`	Materialise a (possibly tensor-parallel) tensor as a plain local tensor.
`default_eagle3_aux_layer_ids`	Return the EAGLE-3 default 3-layer (low / mid / high) aux capture recipe.
`validate_eagle3_aux_layer_ids`	Validate an aux-layer selection against a target of `num_layers` depth.

API

class nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch(
    aux_hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    logits: torch.Tensor | None = None,
    target_probs: torch.Tensor | None = None,
    position_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
)

Dataclass

Target-model supervision for one draft-training batch.

Carries exactly one supervision encoding (validated in __post_init__), both consumed directly by Eagle3TrainerModule.forward:

logits — the target’s full-vocab logits; the draft-vocab projection happens trainer-side. Used by the co-located backend, where the tensor never leaves the GPU.
target_probs + position_mask — the already-projected draft-vocab distribution, so a backend that computes it itself (e.g. a remote server) only transfers draft-vocab-sized tensors.

attention_mask

Tensor

aux_hidden_states

Tensor

doc_remaining

Tensor | None = None

input_ids

Tensor

logits

Tensor | None = None

loss_mask

Tensor

position_ids

Tensor | None = None

position_mask

Tensor | None = None

seq_lens

Tensor | None = None

target_probs

Tensor | None = None

nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch.__post_init__() -> None

nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch.to_trainer_inputs() -> dict[str, torch.Tensor]

Return kwargs for Eagle3TrainerModule.forward, dispatching on whichever supervision encoding this batch carries.

class nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel(
    model: torch.nn.Module,
    aux_layer_ids: typing.Sequence[int] | None = None,
    cp_mesh = None
)

Bases: Eagle3TargetBackend

Co-located backend that captures three auxiliary hidden states from a causal LM.

_cp_size

= cp_mesh.size() if cp_mesh is not None else 1

aux_layer_ids

= self._validate_aux_layer_ids(candidate_ids)

model

= model.eval()

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._check_captured(
    captured: dict[int, torch.Tensor]
) -> None

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._default_aux_layer_ids() -> list[int]

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._get_transformer_layers() -> list[torch.nn.Module]

Return decoder layers as an ordered list indexable by integer.

Supports both the HuggingFace layouts (where layers is a ModuleList) and AutoModel’s custom-impl layouts (where layers is a ModuleDict keyed by str(i)). Returning a plain list normalizes the access pattern for downstream register_forward_hook calls.

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._num_hidden_layers() -> int

Return the target’s decoder depth.

Multimodal targets (e.g. Gemma4ForConditionalGeneration) carry no top-level num_hidden_layers; the text backbone’s depth lives on config.text_config. Text-only targets (and the lightweight config stubs used in tests) expose num_hidden_layers directly.

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._validate_aux_layer_ids(
    aux_layer_ids: typing.Sequence[int]
) -> list[int]

Validate aux-layer selection before any forward hooks are registered.

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch

Run the target model and capture aux hidden states plus logits.

With seq_lens (packing), the target runs with a [B, 1, T, T] block-causal mask and per-document position_ids so its outputs respect document boundaries; the packing metadata is forwarded unshifted to the trainer. seq_lens=None keeps the original 2D-mask path.

nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.get_input_embeddings() -> torch.nn.Embedding

Return the target model input embeddings.

nemo_automodel.components.speculative.eagle.target._shift_left_with_zero(
    tensor: torch.Tensor
) -> torch.Tensor

Shift a batched sequence tensor left and zero-fill the tail.

This matches the reference EAGLE-3 target preparation used by SpecForge: sequence-aligned tensors are shifted with padding(..., left=False). See SpecForge eagle3_target_model.py around the target preparation logic referenced by the user.

nemo_automodel.components.speculative.eagle.target._to_full_tensor(
    tensor: torch.Tensor
) -> torch.Tensor

Materialise a (possibly tensor-parallel) tensor as a plain local tensor.

With a tensor-parallel target the lm_head is column-parallel, so its logits come back as a vocab-sharded DTensor. The draft consumes plain tensors, so gather the full tensor before handing it on. A no-op for an already-plain (unsharded or pure-FSDP-replicated) tensor.

nemo_automodel.components.speculative.eagle.target.default_eagle3_aux_layer_ids(
    num_layers: int
) -> list[int]

Return the EAGLE-3 default 3-layer (low / mid / high) aux capture recipe.

The downstream draft model’s fc projection is sized for exactly num_aux_hidden_states layers (default 3) of concatenated target hidden states. Silently deduplicating collisions on shallow targets would yield fewer than 3 captured tensors and crash later inside the draft fc with a confusing shape-mismatch error — raise here instead so the caller picks 3 distinct in-bounds ids that match the draft config. Shared by every target backend (co-located, remote, SGLang) so they all default identically.

nemo_automodel.components.speculative.eagle.target.validate_eagle3_aux_layer_ids(
    aux_layer_ids: typing.Sequence[int],
    num_layers: int
) -> list[int]

Validate an aux-layer selection against a target of num_layers depth.

Shared by every target backend so an explicit aux_layer_ids is checked identically whether the target runs co-located, remote, or under SGLang.