nemo_automodel.components.speculative.eagle.peagle_trainer

Bases: Module

Draft-side P-EAGLE (parallel-drafting EAGLE-3) trainer module.

Faithful port of speculators’ P-EAGLE (https://github.com/vllm-project/speculators/pull/480): the draft predicts all num_depths tokens in a single parallel forward over a flat, COD-subsampled sequence — it does NOT run EAGLE-3’s autoregressive TTT recurrence.

Per training sequence:

1. COD sampling (:func:generate_cod_sample_indices) draws (anchor_pos, depth): depth 0 keeps every position, depth d keeps a geometrically decaying down_sample_ratio**d fraction. The reference position of each element is anchor_pos + depth.
2. Flat input assembly. All depths are concatenated into one [1, total_sampled] sequence. Depth-0 slots take the real token id and the fc-projected target aux hidden state; depth >= 1 slots take the masked mask_token_id and the single learnable mask_hidden placeholder (projected through the same fc).
3. COD flex attention. A single flex_attention forward with the :func:create_peagle_mask_mod block mask: each element attends to the causal depth-0 context of its document plus earlier-or-equal depths of its own rollout. This is exactly what vLLM’s parallel-drafting runtime sees at inference.
4. Count-normalized KL loss. KL(target || draft) over the draft vocab at every supervised sampled position, normalized by a single total token count — deeper depths (fewer COD positions) naturally contribute less gradient. No 0.8**d schedule.

Batches with batch_size > 1 are processed row-by-row (speculators is batch-size-1); per-row losses are accumulated with a shared denominator so the normalization stays count-based across the whole batch.

Sequence partitioning (sequence_partitions > 1). The flat COD forward attends over n * sum(r**d) positions, so its peak attention / activation memory grows with the context length and OOMs on long sequences. P-EAGLE’s Algorithm 1 (arXiv:2602.01469) splits each sequence into S segments by dependency lineage (:func:assign_cod_segments) and runs a separate forward+backward per segment so only one segment’s activations are resident at a time. The partition is exact: each segment additionally reads every depth-0 position up to its right boundary as key/value context (causal completion), so a segment’s queries see exactly the key/value set they would in the single flat forward — the gradients accumulated across segments equal the single-forward gradient.

The split is caller-driven so the gradient sync stays correct under DDP: :meth:build_peagle_plan (no-grad) assigns COD elements to segments, then the recipe runs one forward(..., peagle_segment=(plan, i)) per segment and owns the backward(). Doing the per-segment backward here (inside a single forward) would bypass DistributedDataParallel’s reducer — its grad all-reduce hooks only fire for backwards over the tensor DDP.forward returns — and silently desynchronize ranks. sequence_partitions == 1 and eval take the single flat forward unchanged.

nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule._forward_peagle_segment(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    peagle_segment: tuple
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics

nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule._peagle_position_loss(
    input_ids_row: torch.Tensor,
    aux_row: torch.Tensor,
    target_logits_row: torch.Tensor,
    anchor_pos: torch.Tensor,
    depth: torch.Tensor,
    orig_positions: torch.Tensor,
    loss_positions: torch.Tensor,
    row_length: torch.Tensor,
    seq_len: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule.build_peagle_plan(
    loss_mask: torch.Tensor
) -> '_PeaglePlan'

nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    aux_hidden_states: torch.Tensor,
    target_logits: torch.Tensor,
    peagle_segment: tuple | None = None
) -> nemo_automodel.components.speculative.eagle.core.Eagle3StepMetrics

Sequence-partitioning plan: one unit per non-empty (row, segment).

A plain (non-container) class on purpose: DistributedDataParallel only scatters tensors and built-in containers across its inputs, so passing this object through forward(peagle_segment=(plan, i)) leaves its tensors intact (a dict/tuple of tensors would be sliced along dim 0 by DDP’s scatter).

Per-position KL(target || draft) over the draft vocabulary.

Matches speculators’ kl_div_loss: log_softmax the draft logits, softmax the target logits, and sum the elementwise KL over the vocab axis. Shapes [*, draft_vocab] -> [*].