`nemo_rl.models.automodel.draft.eagle3_llama`#

EAGLE3 TTT training forward, built on automodel’s LlamaEagle3DraftModel.

Module Contents#

Classes#

`_Eagle3LlamaDecoderLayerWithResidualToggle`	Restores the `norm_before_residual` toggle automodel’s layer lacks.
`Eagle3ForwardTerms`	Per-TTT-step loss numerators/denominators and accuracy counts.
`Eagle3DraftModel`	`LlamaEagle3DraftModel` with an RL-shaped TTT `forward`.

Functions#

`_kl_div_per_position_chunk`	Per-position KL over the draft vocab for one sequence chunk.
`_build_attention_mask`	Dense additive mask: causal AND same-document, padding excluded.

Data#

`_LOSS_REDUCTION_EPS`
`_KL_CHUNK_TOKENS`

API#

nemo_rl.models.automodel.draft.eagle3_llama._LOSS_REDUCTION_EPS#: 1e-05

nemo_rl.models.automodel.draft.eagle3_llama._KL_CHUNK_TOKENS#: 512

class nemo_rl.models.automodel.draft.eagle3_llama._Eagle3LlamaDecoderLayerWithResidualToggle(config, layer_id: int = 0)#

Bases: nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaDecoderLayer

Restores the norm_before_residual toggle automodel’s layer lacks.

Eagle3LlamaDecoderLayer.forward always uses the PRE-normalization hidden state as the residual (residual = hidden_states before hidden_norm runs) – equivalent to norm_before_residual=False, with no toggle at all. Speculators-format checkpoints (e.g. RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3) explicitly declare norm_before_residual: true; loading them through the unmodified automodel layer would silently use the wrong residual convention. Only forward differs from the parent; state dict keys and shapes are unchanged.

Initialization

forward( input_embeds: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor, cache_hidden: list[list[torch.Tensor]], cu_seqlens: Optional[torch.Tensor] = None, max_seqlen: Optional[int] = None, ) → torch.Tensor#

class nemo_rl.models.automodel.draft.eagle3_llama.Eagle3ForwardTerms#

Per-TTT-step loss numerators/denominators and accuracy counts.

Numerators carry gradients; denominators are detached scalars. The caller normalizes as sum_k decay^k * num_k / den_k with den_k optionally reduced across the data-parallel group.

loss_nums: list[torch.Tensor]#: None

loss_dens: list[torch.Tensor]#: None

full_acc_nums: list[torch.Tensor]#: None

full_acc_dens: list[torch.Tensor]#: None

cond_acc_nums: list[torch.Tensor]#: None

cond_acc_dens: list[torch.Tensor]#: None

nemo_rl.models.automodel.draft.eagle3_llama._kl_div_per_position_chunk( logits: torch.Tensor, targets: torch.Tensor, ) → torch.Tensor#

Per-position KL over the draft vocab for one sequence chunk.

Each position’s KL is independent, so calling this per-chunk (as the TTT loop does, alongside chunking compute_logits itself) is mathematically exact and caps the per-chunk TRANSIENT peak – the fp32 [*, chunk, V] intermediates (log_softmax, softmax, elementwise kl) never materialize for the FULL row at once. It does NOT reduce the cumulative memory retained for backward across the whole step: none of the chunks are checkpointed, so autograd still holds onto each chunk’s saved activations until the eventual .backward().

nemo_rl.models.automodel.draft.eagle3_llama._build_attention_mask( document_ids: torch.Tensor, total_seq_len: int, dtype: torch.dtype, device: torch.device, ) → torch.Tensor#

Dense additive mask: causal AND same-document, padding excluded.

Used only for the TTT loop’s step-0 causal block (Eagle3LlamaAttention keeps this mask constant across TTT steps – later steps attend only to their own previous-step K/V at the same position, which needs no mask). documents_ids == -1 marks padding; queries in padding attend nowhere.

class nemo_rl.models.automodel.draft.eagle3_llama.Eagle3DraftModel(config)#

Bases: nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel

LlamaEagle3DraftModel with an RL-shaped TTT forward.

Everything except forward (embedding/fc/lm_head, d2t/t2d, set_vocab_mapping, copy_embeddings_from_target) is inherited unchanged from automodel. model.layers[0] is swapped for

Class:: _Eagle3LlamaDecoderLayerWithResidualToggle right after construction (before any weights load) – same parameter names/shapes, so checkpoints round-trip unchanged.

forward runs the WHOLE TTT unroll as a single top-level call (matching the parent’s calling convention) rather than driving embed_input_ids / project_hidden_states / compute_logits as separate calls: FSDP2’s fully_shard unshard/reshard hooks fire on nn.Module.__call__, and only the top-level model and each decoder layer are individually fully_shard-ed (see build_eagle3_draft_model) – embed_tokens/fc/lm_head/ norm are NOT, so calling their owning methods directly (bypassing draft_model(...)) would run them against still-sharded DTensor parameters.

Initialization

set_embedding_head_trainable(trainable: bool) → None#

get_d2t_target_ids() → torch.Tensor#: Absolute target token id for each draft index (offset convention).

forward( fused_hidden_states: torch.Tensor, input_ids: torch.Tensor, document_ids: torch.Tensor, loss_mask: torch.Tensor, teacher_logits: torch.Tensor, position_ids: Optional[torch.Tensor] = None, ttt_steps: int = 3, seq_lens: Optional[torch.Tensor] = None, ) → nemo_rl.models.automodel.draft.eagle3_llama.Eagle3ForwardTerms#

TTT training forward over one packed row.

Parameters:

fused_hidden_states – [1, T, num_aux * H_target] concatenated aux hiddens.
input_ids – [1, T] rollout token ids (target vocab).
document_ids – [1, T] per-position document index, -1 for padding.
loss_mask – [1, T] boolean positions that carry loss.
teacher_logits – [1, T, draft_vocab] policy logits pre-mapped into draft-vocab order (detached by the caller).
position_ids – [1, T]; defaults to the speculators 1 + arange.
ttt_steps – unroll depth.
seq_lens – [1, batch_size] per-sample block width (each entry the fixed seq_len every sample’s slot occupies in the packed row, INCLUDING that sample’s trailing padding – see Eagle3Runtime.compute_loss). Required when attn_implementation="flash_attention_2", unused otherwise. Declaring each fixed-size block as one FlashAttention-varlen “document” is safe even though padding sits inside the block: padding is placed at each block’s TAIL, so causal masking already keeps every real token from attending to any padding position (padding comes chronologically after), and padding positions’ own outputs never enter the loss (loss_mask excludes them).

nemo_rl.models.automodel.draft.eagle3_llama#