> 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

Target-model wrapper for minimal EAGLE-3 training.

## Module Contents

### Classes

| Name                                                                                             | Description                                                                      |
| ------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- |
| [`Eagle3TargetBatch`](#nemo_automodel-components-speculative-eagle-target-Eagle3TargetBatch)     | Target-model supervision for one draft-training batch.                           |
| [`HFEagle3TargetModel`](#nemo_automodel-components-speculative-eagle-target-HFEagle3TargetModel) | Co-located backend that captures three auxiliary hidden states from a causal LM. |

### Functions

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

### API

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

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

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

```python
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](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend)

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

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

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

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

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

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

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

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

Return the target model input embeddings.

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

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

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

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