> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.eagle.target_v12

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

## Module Contents

### Classes

| Name                                                                                               | Description                                                          |
| -------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`EagleTargetBatch`](#nemo_automodel-components-speculative-eagle-target_v12-EagleTargetBatch)     | Target-model outputs needed by the EAGLE-1 / EAGLE-2 trainer.        |
| [`HFEagleTargetModel`](#nemo_automodel-components-speculative-eagle-target_v12-HFEagleTargetModel) | Thin wrapper that exposes hidden-state supervision from a causal LM. |

### Functions

| Name                                                                                                     | Description                                                              |
| -------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`_shift_left_with_zero`](#nemo_automodel-components-speculative-eagle-target_v12-_shift_left_with_zero) | Shift a batched sequence tensor left and zero-fill the tail.             |
| [`_to_full_tensor`](#nemo_automodel-components-speculative-eagle-target_v12-_to_full_tensor)             | Materialise a (possibly tensor-parallel) tensor as a plain local tensor. |

### API

```python
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.

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

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

```python
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.

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

Return the target model input embeddings.

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

Return the target model lm\_head.

```python
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.

```python
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.