> This page is for version Latest.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.eagle4

EAGLE-4 draft training on top of the EAGLE-3 recipe.

An EAGLE-4 draft is the EAGLE-3.1 Llama-style draft with a sliding-window
draft layer and a low-rank candidate head (see `draft_llama.py`). It trains
in two stages, both run by `TrainEagle3Recipe`:

1. **Draft stage** (`recipe_args.eagle4.sliding_window` only): the usual
   EAGLE-3 TTT loss over the full target vocabulary, with the windowed layer.
2. **Candidate-head stage** (`recipe_args.eagle4.candidate_head` set and
   `draft_weights_path` pointing at the stage-1 export): every stage-1
   parameter is frozen and only the two rank-`r` factors train, initialized
   from the truncated SVD of the draft `lm_head`.
   `Eagle4CandidateHeadTrainerModule` runs the frozen draft's TTT unroll
   and minimizes, at every depth, the forward KL from the full head's softmax
   to the low-rank head's softmax. The low-rank head only has to retrieve the
   full head's likely tokens (serving rescores the retrieved candidates
   exactly), and forward KL weights each token by the full head's probability,
   so the limited rank goes to the tokens that matter. The logged `accuracy`
   is candidate recall: how often the full head's top-1 token is among the
   low-rank head's top `num_candidates`.

`python -m nemo_automodel.components.speculative.eagle.export_eagle4` turns
the stage-2 checkpoint into the `LlamaForCausalLMEagle4` layout the serving
engines load.

## Module Contents

### Classes

| Name                                                                                                                       | Description                                                                   |
| -------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`Eagle4CandidateHeadConfig`](#nemo_automodel-components-speculative-eagle-eagle4-Eagle4CandidateHeadConfig)               | Shape of the EAGLE-4 low-rank candidate head.                                 |
| [`Eagle4CandidateHeadTrainerModule`](#nemo_automodel-components-speculative-eagle-eagle4-Eagle4CandidateHeadTrainerModule) | Stage-2 trainer: fit the low-rank candidate head of a frozen EAGLE-4 draft.   |
| [`Eagle4Config`](#nemo_automodel-components-speculative-eagle-eagle4-Eagle4Config)                                         | `recipe_args.eagle4`: turns the EAGLE-3 recipe's draft into an EAGLE-4 draft. |

### Data

[`CANDIDATE_HEAD_KEYS`](#nemo_automodel-components-speculative-eagle-eagle4-CANDIDATE_HEAD_KEYS)

### API

```python
class nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadConfig(
    rank: int = 128,
    num_candidates: int = 128
)
```

Dataclass

Shape of the EAGLE-4 low-rank candidate head.

**`num_candidates`** `int = 128`

Tokens `K` the serving engine keeps from the low-rank
scores and rescores with the full `lm_head`; also the `K` of the
logged candidate recall.

---

**`rank`** `int = 128`

Rank `r` of the factorized vocabulary scores.

---

```python
nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadConfig.__post_init__() -> None
```

```python
nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadConfig.build(
    draft_model: nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel,
    ttt_steps: int
) -> nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadTrainerModule
```

Build the stage-2 trainer that fits `draft_model`'s candidate head.

```python
class nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadTrainerModule(
    draft_model: nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel,
    ttt_steps: int,
    num_candidates: int
)
```

**Bases:** `Module`

Stage-2 trainer: fit the low-rank candidate head of a frozen EAGLE-4 draft.

The draft's TTT unroll (same recurrence, input shift and packing gates as
`Eagle3TrainerModule`) runs without gradients; at each depth the full
`lm_head` softmax is the teacher and the low-rank scores the student.
Depth `i` is weighted by `0.8 ** i` and the sum normalized by the total
weight, as in the draft stage.

Like every EAGLE-3-recipe trainer it carries the draft-vocab mapping buffers
(`selected_token_ids` / `selected_token_mask`) that the recipe saves to and
restores from the checkpoint metadata; an EAGLE-4 draft uses the full vocabulary.

```python
nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadTrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_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
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics
```

Run the candidate-head loss for one batch.

Accepts the same inputs as `Eagle3TrainerModule.forward` so every
target backend can feed it. The target supervision (`target_logits`,
or `target_probs` / `position_mask` from the offline cache) is not
used: the teacher is the draft's own full head, and `loss_mask`
selects the supervised positions.

**Parameters:**

**`input_ids`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, token ids.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, right-padded
`1` / `0` validity mask.

---

**`loss_mask`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, `1` where the
draft is supervised.

---

**`aux_hidden_states`** `torch.Tensor`

Tensor of shape `[batch, sequence, 3 * target_hidden]`,
the concatenated target hidden states the draft projects.

---

**`target_logits`** `torch.Tensor | None` — default: None

Unused; tensor of shape `[batch, sequence, vocab]`.

---

**`target_probs`** `torch.Tensor | None` — default: None

Unused; tensor of shape `[batch, sequence, vocab]`.

---

**`position_mask`** `torch.Tensor | None` — default: None

Unused; tensor of shape `[batch, sequence, 1]`.

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional tensor of shape `[batch, sequence]` with
per-document positions (packing).

---

**`seq_lens`** `torch.Tensor | None` — default: None

Optional tensor of shape `[batch, max_docs]` with packed
document lengths, 0-padded.

---

**`doc_remaining`** `torch.Tensor | None` — default: None

Optional tensor of shape `[batch, sequence]`: tokens
left in each slot's document; depth `i` is supervised only
while `i < doc_remaining`.

---

**Returns:** `Eagle3StepMetrics`

`Eagle3StepMetrics` with the weighted forward KL (0 for a perfect head), `accuracy` set to the

```python
class nemo_automodel.components.speculative.eagle.eagle4.Eagle4Config(
    sliding_window: int | None = None,
    candidate_head: nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadConfig | None = None
)
```

Dataclass

`recipe_args.eagle4`: turns the EAGLE-3 recipe's draft into an EAGLE-4 draft.

**`candidate_head`** `Eagle4CandidateHeadConfig | None = None`

Set for the candidate-head stage; `None` trains the
draft itself.

---

**`sliding_window`** `int | None = None`

Keys each draft query attends to, itself and its TTT
chain included (so at least `ttt_steps`), or `None` for full
causal attention.

---

```python
nemo_automodel.components.speculative.eagle.eagle4.Eagle4Config.__post_init__() -> None
```

```python
nemo_automodel.components.speculative.eagle.eagle4.Eagle4Config.draft_config_updates(
    aux_layer_ids: list[int]
) -> dict[str, typing.Any]
```

Return the draft `config.json` keys that define this EAGLE-4 draft.

**Parameters:**

**`aux_layer_ids`** `list[int]`

The three target layers whose outputs feed the draft,
recorded as `eagle_config.eagle_aux_hidden_state_layer_ids` so
the serving engine captures the same layers.

---

```python
nemo_automodel.components.speculative.eagle.eagle4.CANDIDATE_HEAD_KEYS = frozenset({'efficient_lm_head.weight', 'efficient_lm_head_down_proj.weight'})
```