nemo_automodel.components.speculative.eagle.target_v12

Target-model wrapper for EAGLE-1 / EAGLE-2 training.

Module Contents

Classes

Name	Description
`EagleTargetBatch`	Target-model outputs needed by the EAGLE-1 / EAGLE-2 trainer.
`HFEagleTargetModel`	Thin wrapper that exposes hidden-state supervision 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.

API

class nemo_automodel.components.speculative.eagle.target_v12.EagleTargetBatch(
    input_hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    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
)

Dataclass

Target-model outputs needed by the EAGLE-1 / EAGLE-2 trainer.

position_ids / seq_lens / doc_remaining are None on the unpacked path and carry the packing metadata (unshifted, indexed by slot) through to the trainer on the packed path.

attention_mask

Tensor

doc_remaining

Tensor | None = None

input_hidden_states

Tensor

input_ids

Tensor

loss_mask

Tensor

position_ids

Tensor | None = None

seq_lens

Tensor | None = None

target_hidden_states

Tensor

target_logits

Tensor

class nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel(
    model: torch.nn.Module
)

Thin wrapper that exposes hidden-state supervision from a causal LM.

model

= model.eval()

nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.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_v12.EagleTargetBatch

Run the target transformer and prepare shifted supervision tensors.

All per-token inputs are [B, T]. With seq_lens ([B, max_docs] long, per-document lengths summing to T) the target runs with a document-level block-causal mask and per-document position_ids so its hidden states do not leak across document boundaries; SDPA/eager targets consume the [B, 1, T, T] block-causal additive mask, FlashAttention targets infer document boundaries from position_ids and are passed attention_mask=None (batch size 1 only). position_ids / seq_lens / doc_remaining are carried through (unshifted) so the trainer can build the draft’s block-causal mask and drop cross-document supervision.

nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_input_embeddings() -> torch.nn.Embedding

Return the target model input embeddings.

nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_lm_head() -> torch.nn.Module

Return the target model lm_head.

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

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

nemo_automodel.components.speculative.eagle.target_v12._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.