nemo_automodel.components.speculative.eagle.core

Aggregated metrics from one EAGLE-3 training step.

step_prefix_hits / step_valid are [ttt_steps] per-TTT-step counts for the simulated accept length (:func:simulated_accept_length): step_prefix_hits[k] counts positions whose greedy draft chain is still fully correct through depth k + 1 (top-1 hit at every step <= k), and step_valid[k] counts positions supervised at every step <= k under the shifted loss/document masks, so numerator and denominator cover the same chain population (a chain with an unsupervised earlier depth, e.g. from a gappy multi-turn loss mask, can never be a hit and must not count as a miss). step_valid deliberately does NOT exclude positions whose target token falls outside the compressed draft vocabulary: serving must reject those tokens (the draft cannot emit them), so they count as chain breaks instead of dropping out of the denominator. Both are kept unreduced so the recipe can accumulate them over a logging window. The P-EAGLE trainer, whose depths are drafted in parallel from one anchor rather than as a TTT chain, leaves them None.

Bases: Module

Draft-side EAGLE-3 trainer module with test-time-training unroll.

nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule._lk_step_loss(
    logits: torch.Tensor,
    target_probs: torch.Tensor,
    position_mask: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule.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

Bases: Function

Differentiable sum-all-reduce across the cp group.

Forward sums per-rank inputs into a replicated total. Each rank uses that total identically, so the loss gradient w.r.t. this rank’s input is just the incoming grad (coefficient 1) — hence the identity backward.

nemo_automodel.components.speculative.eagle.core._CpAllReduceSum.backward(
    ctx,
    grad
)

nemo_automodel.components.speculative.eagle.core._CpAllReduceSum.forward(
    ctx,
    x,
    cp_group
)

Project target logits into draft vocabulary space and build supervision mask.

Renormalize a per-shard masked-mean loss over the full cp sequence.

step_loss is the mean over this rank’s LOCAL supervised positions. Recover the local sum (mean * count), sum both across cp (the sum differentiably), and divide by the global count — so the loss value equals the full-sequence loss and the backprop’d gradient is the global gradient.

Left-shift a context-parallel-sharded sequence tensor by one position.

Each rank holds a contiguous S/cp shard. A plain left-shift would zero-fill the boundary, but the token that rolls into rank r’s tail lives at rank r+1’s head. Shift locally, then P2P the neighbour’s (current) head into the tail; the last rank keeps the zero fill. Applied in lockstep each TTT step this reproduces the global shift exactly. Used for labels / input_ids only (no grad).

Left-shift a ZIG-ZAG-sharded sequence tensor by one position.

Under zig-zag sharding rank r holds two non-contiguous global chunks — the early chunk r and the late chunk 2*cp-1-r — laid out locally as [early | late] (each of length c = local_len/2). A global left-shift therefore rolls each half locally and needs two boundary tokens:

• the early half’s tail (global end of chunk r) is the head of chunk r+1 — rank r+1’s early head — except on the last rank, where chunk r+1 is that rank’s OWN late chunk, so the fill is local.
• the late half’s tail (global end of chunk 2*cp-1-r) is the head of chunk 2*cp-r — rank r-1’s late head — except on rank 0, whose late tail is the global last position and stays zero-filled.

Two ring P2P exchanges (early head -> r-1, late head -> r+1) cover both. Used for labels / input_ids only (no grad). Applied in lockstep each TTT step this reproduces the global shift exactly. See :func:_cp_shift_left for the contiguous-shard analogue.

Mean of values over the supervised positions.

Parameters:

Returns: torch.Tensor

Scalar tensor; the masked mean (zero when no position is supervised).

Shift a batched sequence tensor left and zero-fill the tail.

Expected accepted tokens per speculative round, from prefix-hit counts.

Models greedy chain drafting: step_prefix_hits[k] / step_valid[k] estimates the joint probability that the first k + 1 drafted tokens all match the target’s greedy choices, i.e. that the chain survives depth k + 1, so the expectation is 1 + sum_k P(survives depth k + 1). Joint prefix counts keep the correlation between depths that a product of per-step marginal accuracies discards (coincident and disjoint hits score differently). The leading 1 counts the token the target itself emits on every verification round, matching the accept_length convention of the serving benchmarks (1 + accepted/drafts). A step with no supervised positions contributes zero via the clamp_min. This is a training-time proxy for greedy chain decoding; engine tree drafting typically accepts more.