Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsSpeculativeEaglenemo_automodel.components.speculative.eagle.eagle4

nemo_automodel.components.speculative.eagle.eagle4

View as Markdown

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

NameDescription
Eagle4CandidateHeadConfigShape of the EAGLE-4 low-rank candidate head.
Eagle4CandidateHeadTrainerModuleStage-2 trainer: fit the low-rank candidate head of a frozen EAGLE-4 draft.
Eagle4Configrecipe_args.eagle4: turns the EAGLE-3 recipe’s draft into an EAGLE-4 draft.

Data

CANDIDATE_HEAD_KEYS

API

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.

nemo_automodel.components.speculative.eagle.eagle4.Eagle4CandidateHeadConfig.__post_init__() -> None
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.

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.

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 | NoneDefaults to None

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

target_probs
torch.Tensor | NoneDefaults to None

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

position_mask
torch.Tensor | NoneDefaults to None

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

position_ids
torch.Tensor | NoneDefaults to None

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

seq_lens
torch.Tensor | NoneDefaults to None

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

doc_remaining
torch.Tensor | NoneDefaults to 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

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.

nemo_automodel.components.speculative.eagle.eagle4.Eagle4Config.__post_init__() -> None
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.

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