nemo_automodel.components.speculative.eagle

Release backend resources (remote connections, server handles).

Run the target and return the supervision for one training batch.

Submit an asynchronous generate_batch for prefetch pipelining.

Only backends that overlap target inference with draft training implement this; the default signals a synchronous backend so callers fall back to generate_batch.

Return the target input-embedding module (used to seed the draft).

Provide the draft-vocab mapping needed to precompute supervision.

Co-located backends keep the mapping on the trainer module and derive the distribution there, so the default is a no-op. A backend that computes target_probs itself overrides this to receive the mapping.

One TTT step of the LK acceptance-rate loss (arXiv:2602.23881).

The per-token expected acceptance under speculative sampling is the min-overlap of the two distributions, alpha = sum_v min(p_target, p_draft), computed over the draft vocabulary (both inputs already live there).

• "alpha": the pure acceptance likelihood -mean(log alpha); the log is taken per token before the masked mean, with zero-acceptance tokens contributing zero (they carry no usable gradient direction).
• "lambda": the adaptive hybrid kl_weight * soft_ce + (1 - kl_weight) * (1 - mean(alpha)) with kl_weight = lk_kl_scale * exp(-lk_kl_decay * mean(alpha).detach()); the weight is detached so gradients flow only through the two terms, and training shifts from distillation toward direct acceptance as the draft improves. A step with zero supervised positions returns 0, matching the soft-CE path.

Under CP every masked mean is renormalized over the full sequence via _cp_global_step_loss, so the mixing weight and the loss match the single-rank run.

Parameters:

Returns: torch.Tensor

Scalar loss tensor for this TTT step.

Run the EAGLE-3 unrolled draft loss for one batch.

The attention layer is driven through a shared cache_hidden list so each TTT step can attend to the K/V branches produced by every previous step at the same position. This matches the SpecForge llama3_eagle.py recurrence; without it, multi-step TTT would degenerate into ttt_steps independent single-step passes and the draft would never learn the multi-step distribution it sees at deployment time.

attention_mask is held constant across TTT steps — only input_ids / loss_mask / position_mask / target_probs roll forward by one position per step.

Packing: position_ids / seq_lens make the draft’s Block-1 attention document-level block-causal, and doc_remaining gates supervision per step (slot t valid at step k only while k < doc_remaining[t]), masking every cross-document TTT prediction.

Two supervision sources are accepted: the live path passes the target’s full-vocab target_logits and the draft distribution is derived here; the offline-cache path (precompute_eagle3) passes the already-derived target_probs (over the draft vocab) and position_mask directly, so the full-vocab logits never have to be stored. Provide exactly one of the two.

Project predicted hidden states through the frozen target lm_head.

Run one EAGLE-1 / EAGLE-2 training step.

Per-token tensors are [B, T] (position_ids / doc_remaining included) except the [B, T, H] hidden states and [B, T, V] target_logits; seq_lens is [B, max_docs]. When packing is on, position_ids / seq_lens make the draft block-causal and per-document, and doc_remaining (real tokens after each slot within its document) gates supervision: a document’s last real token (doc_remaining == 0) is dropped because the wrapper’s global left-shift makes its target the next document’s first token.

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.

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.

Validate aux-layer selection before any forward hooks are registered.

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.

Return the target model input embeddings.

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.

Return the target model input embeddings.

Return the target model lm_head.

Run a VLM and construct shifted multimodal draft supervision.

input_ids and attention_mask are read from model_inputs rather than taken separately, so the tensors the target actually runs on are the same ones the image mask and the returned attention mask are derived from. A second copy would misalign the draft’s modality routing without raising.

Parameters:

Returns: MSDTargetBatch

MSDTargetBatch with shifted image-aware embeddings, target hidden

Return the frozen VLM language-model head used by the draft loss.

Compute draft logits on the configured draft vocabulary.

With config.norm_output unset (EAGLE-3 default) the input is the raw decoder-layer output and the final model.norm is applied here. With config.norm_output set (EAGLE-3.1) forward has already returned the post-norm state, so lm_head is applied directly to avoid a double normalisation.

Initialize draft embeddings from the target model embeddings.

When the target model is wrapped with FSDP2, target_embedding.weight is a DTensor sharded across ranks. The draft embedding is a plain nn.Parameter (the draft is not FSDP-wrapped), so a direct copy_ of a DTensor into a regular tensor raises a mixed-type distributed-operator error. Gather to a full local tensor first.

Embed input ids with the draft embedding table.

Run one full-sequence draft update step.

cache_hidden is the EAGLE-3 TTT cache. Pass [[], []] on the first step of a TTT unroll and the same list object on each subsequent step; the attention layer appends the per-step K and V to it. If None is passed (e.g. from a one-shot evaluation call) a fresh [[], []] is allocated locally — step 0 of TTT is mathematically equivalent to a plain causal forward.

seq_lens (packing) makes Block-1 attention document-level block-causal (eager mask / FA2 varlen); callers must pass per-document position_ids.

Freeze draft input embeddings.

Disable activation checkpointing for the P-EAGLE draft layers.

Enable activation checkpointing for the P-EAGLE draft layers.

Training-only memory knob: recomputes each forward_peagle layer in the backward instead of storing its activations (the EAGLE-3 TTT forward path is unaffected). gradient_checkpointing_kwargs is accepted for HF-API parity but ignored — recompute is always non-reentrant, the only mode compatible with the non-tensor block_mask.

Project concatenated target aux states from num_aux * H_target to draft hidden size.

When config.fc_norm is set (EAGLE-3.1), the input is split into num_aux_hidden_states equal chunks along the last dim and each chunk is passed through its own RMSNorm in model.fc_norm (the modules are independent, matching vLLM’s upstream implementation). The normalized chunks are then re-concatenated and fed to fc, stabilising the per-aux-state scale before the projection mixes them and removing the speculation-depth drift observed with raw inputs.

Populate the d2t / t2d vocab-remap buffers from the draft->target id map.

selected_token_ids has shape [draft_vocab_size]; entry i is the target vocab id of draft id i (the frequency-pruned mapping built by build_eagle3_token_mapping). This writes the two buffers inference engines consume:

• d2t[i] = selected_token_ids[i] - i — the offset form vLLM expects (target_id = draft_id + d2t[draft_id]);
• t2d[target_id] = True for every selected target id — the boolean presence mask SGLang consumes.

These must be in the saved checkpoint: without them vLLM/SGLang find no mapping, silently align draft ids to the first draft_vocab_size target ids, and acceptance rate collapses.

No-op when the draft vocab is not compressed (the buffers do not exist and the draft logits are already in target space).

Copy the target model token embeddings into the draft embeddings.

When the target is wrapped with FSDP2, its embed_tokens.weight is a DTensor sharded across ranks. Gather to a local full tensor before copying into the (unsharded) draft parameter — otherwise aten.copy_ raises a mixed Tensor/DTensor error.

Predict the next-step target hidden state for each position.

Args (batch B, sequence length T, hidden size H): input_ids: [B, T] long token ids. target_hidden_states: [B, T, H] target hidden features fed to the draft. attention_mask: [B, T] (1 = real, 0 = pad); used on the unpacked path only. position_ids: [B, T] long, or None to default to arange. Packing requires per-document positions (reset to range(doc_len)). seq_lens: [B, max_docs] long (0-padded per-document lengths that sum to T per row) turns on sequence packing — attention becomes document-level block-causal via build_block_causal_additive_mask instead of the plain causal + padding mask, and RoPE uses the per-document position_ids so each document is rotated from its own origin.

Freeze draft token embeddings.

Restore curriculum progress from a checkpoint state dictionary.

Return the multimodal batch probability at global_step.

Return the current source selection and advance one optimizer step.

Return the curriculum phase active at global_step.

Deterministically choose the data source for global_step.

The random draw is seeded from the global step rather than mutable RNG state. Resumed and distributed runs therefore make the same sampling decision for a given optimizer step.

Return progress required to resume next_source.

Find the language token embedding module across common HF VLM layouts.

Draft and greedily verify one MSD candidate tree for a VLM prompt.

Project predicted hidden states through the frozen target language head.

Parameters:

Returns: torch.Tensor

Tensor of shape […, vocab], containing target-vocabulary logits.

Run one multimodal speculative-decoding training step.

Parameters:

Returns: MSDStepMetrics

MSDStepMetrics containing scalar loss terms, first-token accuracy, and

Create a feature-level candidate tree after a target-sampled root.

shifted_inputs_embeds and shifted_image_mask use the same next-token alignment as HFMSDTargetModel. The final shifted embedding is replaced with the target root token embedding, then each recursive draft prediction is appended as the feature context for its child nodes. attention_mask must be right-padded, since the prefix is taken as its leading active slice.

Return the root-to-leaf candidate-token path.

Return all leaf paths in deterministic tree order.

Raise: MSD drafts consume target embeddings instead of holding their own.

Predict next-position target hidden states from multimodal features.

Parameters:

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden], the draft predictions for

Raise: MSD drafts have no token embedding table to freeze.

Loss for one segment of a build_peagle_plan plan.

The recipe drives this once per plan.units entry and back-propagates each result, so each segment owns a self-contained autograd graph that is freed before the next — and the backward flows through DDP.forward so gradients all-reduce correctly. metrics.loss is the segment’s share of the count-normalized batch loss (loss / total_den); summing it over the plan reproduces the single flat forward’s loss.

Draft forward + count-normalized KL for one row’s COD elements.

Shared by the single flat forward (one call per row, all sampled positions charged) and the partitioned forward (one call per segment, only the segment’s owned/supervised positions charged — the rest ride along as key/value context). Returns (loss_num, loss_den, correct, valid) as float scalars, where the loss is Σ KL over loss_positions and loss_den is their count; the caller normalizes.

Assign COD elements to segments for the sequence-partitioning path.

Samples COD once per row (the indices must be reused across the segment forwards), runs Algorithm 1 assignment (assign_cod_segments) plus causal completion, and emits one unit per non-empty (row, segment) as (b, anchor, depth, orig_positions, loss_positions). loss_positions marks the segment’s owned supervised elements (charged loss); the other elements are depth-0 causal-completion context (key/value only). The shared total_den is the batch’s total supervised-token count, so each segment’s loss / total_den sums to the single-forward loss.

Run the P-EAGLE parallel-drafting loss for one batch.

attention_mask supplies the per-row valid length so padded positions are excluded from attention (document mask) and from supervision.

peagle_segment selects the sequence-partitioning path: when it is a (plan, index) pair (built by build_peagle_plan) this computes the loss for that one segment only — the recipe calls this once per segment and owns the backward() so DDP’s gradient sync stays correct. When None (sequence_partitions == 1 and eval) a single flat forward over the whole COD sequence returns a grad-carrying loss.

Release the runner (frees GPU memory / engine handles).

Run the runner’s target and capture aux hidden states plus logits.

Produces an Eagle3TargetBatch byte-for-byte compatible with HFEagle3TargetModel.generate_batch: the logits, input_ids and loss_mask are shifted left by one (next-token alignment) while the aux hidden states are kept position-aligned. The draft-vocab projection happens trainer-side / server-side from logits, exactly as in the co-located path.

Note: sequences are assumed right-padded (loss_mask zeros the pad). With causal attention, trailing pad tokens do not affect earlier positions, so the captured supervision matches a masked HuggingFace forward.

Return an object exposing .weight (the target input embeddings).

Matches the offline-cache / remote path: the draft’s copy_embeddings_from_target only reads .weight.

Run the target once and return (logits, aux_hidden_states).

logits is [batch, seq, vocab] (full vocab, unshifted) and aux_hidden_states is [batch, seq, 3 * hidden] (the three capture layers concatenated on the last dim, unshifted).

Return the target input-embedding weight [vocab, hidden].

Tell the underlying model which 3 decoder layers to capture.