> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.speculative.eagle EAGLE-3 training components. Only the EAGLE-3-native model and trainer classes live here. Generic helpers were moved to the project's canonical locations: * `masked_soft_cross_entropy` -> `nemo_automodel.components.loss.soft_ce` * `build_eagle3_dataloader` / `build_eagle3_token_mapping` -> `nemo_automodel.components.datasets.llm.eagle3` Import them directly from those modules. ## Subpackages * **[`nemo_automodel.components.speculative.eagle.remote`](/nemo-automodel/nemo_automodel/components/speculative/eagle/remote)** ## Submodules * **[`nemo_automodel.components.speculative.eagle.backend`](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend)** * **[`nemo_automodel.components.speculative.eagle.core`](/nemo-automodel/nemo_automodel/components/speculative/eagle/core)** * **[`nemo_automodel.components.speculative.eagle.core_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/core_v12)** * **[`nemo_automodel.components.speculative.eagle.draft_deepseek`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_deepseek)** * **[`nemo_automodel.components.speculative.eagle.draft_gemma`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gemma)** * **[`nemo_automodel.components.speculative.eagle.draft_gpt_oss`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gpt_oss)** * **[`nemo_automodel.components.speculative.eagle.draft_kimi_k3`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_kimi_k3)** * **[`nemo_automodel.components.speculative.eagle.draft_llama`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama)** * **[`nemo_automodel.components.speculative.eagle.draft_llama_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12)** * **[`nemo_automodel.components.speculative.eagle.msd`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd)** * **[`nemo_automodel.components.speculative.eagle.msd_curriculum`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_curriculum)** * **[`nemo_automodel.components.speculative.eagle.msd_decode`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_decode)** * **[`nemo_automodel.components.speculative.eagle.msd_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_target)** * **[`nemo_automodel.components.speculative.eagle.peagle_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_attention)** * **[`nemo_automodel.components.speculative.eagle.peagle_data`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_data)** * **[`nemo_automodel.components.speculative.eagle.peagle_draft`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_draft)** * **[`nemo_automodel.components.speculative.eagle.peagle_trainer`](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_trainer)** * **[`nemo_automodel.components.speculative.eagle.registry`](/nemo-automodel/nemo_automodel/components/speculative/eagle/registry)** * **[`nemo_automodel.components.speculative.eagle.ring_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/ring_attention)** * **[`nemo_automodel.components.speculative.eagle.sglang_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_runner)** * **[`nemo_automodel.components.speculative.eagle.sglang_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_target)** * **[`nemo_automodel.components.speculative.eagle.target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target)** * **[`nemo_automodel.components.speculative.eagle.target_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target_runner)** * **[`nemo_automodel.components.speculative.eagle.target_v12`](/nemo-automodel/nemo_automodel/components/speculative/eagle/target_v12)** * **[`nemo_automodel.components.speculative.eagle.ulysses_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/ulysses_attention)** * **[`nemo_automodel.components.speculative.eagle.vispec_core`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_core)** * **[`nemo_automodel.components.speculative.eagle.vispec_draft`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_draft)** * **[`nemo_automodel.components.speculative.eagle.vispec_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_target)** * **[`nemo_automodel.components.speculative.eagle.vllm_runner`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_runner)** * **[`nemo_automodel.components.speculative.eagle.vllm_target`](/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_target)** * **[`nemo_automodel.components.speculative.eagle.zigzag_ring_attention`](/nemo-automodel/nemo_automodel/components/speculative/eagle/zigzag_ring_attention)** ## Package Contents ### Classes | Name | Description | | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- | | [`DraftSpec`](#nemo_automodel-components-speculative-eagle-registry-DraftSpec) | How to build an EAGLE draft model for a particular target architecture. | | [`Eagle3TargetBackend`](#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend) | Abstract contract every EAGLE-3 target-model backend implements. | | [`Eagle3TrainerModule`](#nemo_automodel-components-speculative-eagle-core-Eagle3TrainerModule) | Draft-side EAGLE-3 trainer module with test-time-training unroll. | | [`EagleTrainerModule`](#nemo_automodel-components-speculative-eagle-core_v12-EagleTrainerModule) | Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction. | | [`GptOssEagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_gpt_oss-GptOssEagle3DraftModel) | EAGLE-3 draft model for gpt-oss targets. | | [`HFEagle3TargetModel`](#nemo_automodel-components-speculative-eagle-target-HFEagle3TargetModel) | Co-located backend that captures three auxiliary hidden states from a causal LM. | | [`HFEagleTargetModel`](#nemo_automodel-components-speculative-eagle-target_v12-HFEagleTargetModel) | Thin wrapper that exposes hidden-state supervision from a causal LM. | | [`HFMSDTargetModel`](#nemo_automodel-components-speculative-eagle-msd_target-HFMSDTargetModel) | Expose image-aware feature supervision from a frozen Hugging Face VLM. | | [`LlamaEagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_llama-LlamaEagle3DraftModel) | Llama-style dense EAGLE-3 draft model (Llama, Phi-3, Qwen3). | | [`LlamaEagleDraftModel`](#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel) | Llama-style dense draft that predicts next-step hidden states. | | [`MSDCurriculum`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculum) | Stateless-by-step two-stage curriculum with checkpointable progress. | | [`MSDCurriculumConfig`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculumConfig) | Configuration for the two-stage MSD data curriculum. | | [`MSDCurriculumPhase`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDCurriculumPhase) | High-level phase of the two-stage MSD training schedule. | | [`MSDDataSource`](#nemo_automodel-components-speculative-eagle-msd_curriculum-MSDDataSource) | Data source selected for one MSD optimizer step. | | [`MSDGreedyDecoder`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDGreedyDecoder) | Batch-one reference MSD round with recursive drafting and tree verification. | | [`MSDTrainerModule`](#nemo_automodel-components-speculative-eagle-msd-MSDTrainerModule) | Draft-side training module for multimodal speculative decoding. | | [`MSDTreeDraftGenerator`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeDraftGenerator) | Recursively draft a top-k MSD feature tree from target prefix features. | | [`MSDTreeLayout`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeLayout) | Tree attention and retrieval metadata for an MSD proposal. | | [`MSDTreeNode`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeNode) | One drafted text-token node in an MSD candidate tree. | | [`MSDTreeProposal`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDTreeProposal) | A target root token plus recursively drafted candidate-tree nodes. | | [`MSDVerificationResult`](#nemo_automodel-components-speculative-eagle-msd_decode-MSDVerificationResult) | Greedy posterior verification result for one candidate tree. | | [`MultimodalEagleDraftModel`](#nemo_automodel-components-speculative-eagle-msd-MultimodalEagleDraftModel) | EAGLE-1/2 draft model with modality-aware input feature construction. | | [`PEagleTrainerModule`](#nemo_automodel-components-speculative-eagle-peagle_trainer-PEagleTrainerModule) | Draft-side P-EAGLE (parallel-drafting EAGLE-3) trainer module. | | [`RunnerEagle3TargetModel`](#nemo_automodel-components-speculative-eagle-target_runner-RunnerEagle3TargetModel) | EAGLE-3 target backend that runs the frozen target through a runner. | | [`TargetRunner`](#nemo_automodel-components-speculative-eagle-target_runner-TargetRunner) | Minimal engine surface `RunnerEagle3TargetModel` depends on. | ### Functions | Name | Description | | -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- | | [`accept_or_resample`](#nemo_automodel-components-speculative-eagle-msd_decode-accept_or_resample) | Apply the lossless speculative acceptance rule for one candidate token. | | [`build_msd_tree_layout`](#nemo_automodel-components-speculative-eagle-msd_decode-build_msd_tree_layout) | Build engine-ready tree attention and root-to-leaf retrieval metadata. | | [`resolve_eagle1_draft_spec`](#nemo_automodel-components-speculative-eagle-registry-resolve_eagle1_draft_spec) | Resolve the EAGLE-1 / EAGLE-2 draft spec for a target's `config.architectures` field. | | [`resolve_eagle3_draft_spec`](#nemo_automodel-components-speculative-eagle-registry-resolve_eagle3_draft_spec) | Resolve the EAGLE-3 draft spec for a target's `config.architectures` field. | | [`simulated_accept_length`](#nemo_automodel-components-speculative-eagle-core-simulated_accept_length) | Expected accepted tokens per speculative round, from prefix-hit counts. | | [`verify_greedy_tree`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_greedy_tree) | Select the longest target-greedy accepted path and its bonus token. | | [`verify_hf_greedy_tree`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_hf_greedy_tree) | Reference target verification for a batch-one Hugging Face VLM. | | [`verify_stochastic_chain`](#nemo_automodel-components-speculative-eagle-msd_decode-verify_stochastic_chain) | Losslessly verify a linear draft chain and emit the correction token. | ### Data [`EAGLE1_DRAFT_REGISTRY`](#nemo_automodel-components-speculative-eagle-registry-EAGLE1_DRAFT_REGISTRY) [`EAGLE3_DRAFT_REGISTRY`](#nemo_automodel-components-speculative-eagle-registry-EAGLE3_DRAFT_REGISTRY) ### API ```python class nemo_automodel.components.speculative.eagle.registry.DraftSpec( draft_cls: type[transformers.PreTrainedModel] ) ``` Dataclass How to build an EAGLE draft model for a particular target architecture. **`draft_cls`** `type[PreTrainedModel]` --- ```python class nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend() ``` Abstract Abstract contract every EAGLE-3 target-model backend implements. Two supervision encodings are allowed, both consumed directly by `Eagle3TrainerModule.forward`: * **full logits** -- `Eagle3TargetBatch.logits` carries the target's full-vocab logits and the draft-vocab projection happens trainer-side. Cheap when co-located (the tensor never leaves the GPU); impractical to ship over a wire because it is full-vocab sized. * **precomputed** -- `Eagle3TargetBatch.target_probs` and `Eagle3TargetBatch.position_mask` carry the already-projected draft-vocab distribution, so a backend that computes them itself (e.g. a remote server) only has to transfer draft-vocab-sized tensors. A backend returns exactly one of the two encodings from `generate_batch`; the recipe forwards whichever is present. **`supports_async`** `bool` Whether `generate_batch_async` is implemented (prefetch-capable). --- ```python nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.close() -> None ``` Release backend resources (remote connections, server handles). ```python nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.generate_batch( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor ) -> 'Eagle3TargetBatch' ``` abstract Run the target and return the supervision for one training batch. ```python nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.generate_batch_async( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor ) ``` 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`. ```python nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.get_input_embeddings() -> torch.nn.Module ``` abstract Return the target input-embedding module (used to seed the draft). ```python nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend.set_vocab_mapping( selected_token_ids: torch.Tensor, selected_token_mask: torch.Tensor ) -> None ``` 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. ```python class nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule( draft_model: torch.nn.Module, selected_token_ids: torch.Tensor, selected_token_mask: torch.Tensor, ttt_steps: int, cp_group = None, cp_zigzag: bool = False, lk_loss_type: str | None = None, lk_kl_scale: float = 1.0, lk_kl_decay: float = 3.0 ) ``` **Bases:** `Module` Draft-side EAGLE-3 trainer module with test-time-training unroll. **`lk_kl_decay`** `= float(lk_kl_decay)` --- **`lk_kl_scale`** `= float(lk_kl_scale)` --- ```python nemo_automodel.components.speculative.eagle.core.Eagle3TrainerModule._lk_step_loss( logits: torch.Tensor, target_probs: torch.Tensor, position_mask: torch.Tensor ) -> torch.Tensor ``` 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:** **`logits`** `torch.Tensor` Tensor of shape `[batch, sequence, draft_vocab]`; draft logits. --- **`target_probs`** `torch.Tensor` Tensor of shape `[batch, sequence, draft_vocab]`; target distribution restricted to the draft vocabulary. --- **`position_mask`** `torch.Tensor` Bool tensor of shape `[batch, sequence, 1]`; True where the position is supervised this step. --- **Returns:** `torch.Tensor` Scalar loss tensor for this TTT step. ```python 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 ``` 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. ```python class nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule( draft_model: torch.nn.Module, target_lm_head: torch.nn.Module, hidden_loss_weight: float = 1.0, token_loss_weight: float = 0.1, feature_noise_config: nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig | None = None, rank_loss_weight: float = 0.0, rank_loss_topk: int = 10 ) ``` **Bases:** `Module` Draft-side trainer for EAGLE-1 / EAGLE-2 hidden-state prediction. **`hidden_loss_fn`** `= nn.SmoothL1Loss(reduction='none')` --- ```python nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.compute_logits( hidden_states: torch.Tensor ) -> torch.Tensor ``` Project predicted hidden states through the frozen target lm\_head. ```python nemo_automodel.components.speculative.eagle.core_v12.EagleTrainerModule.forward( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, input_hidden_states: torch.Tensor, target_hidden_states: torch.Tensor, target_logits: torch.Tensor, position_ids: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None, doc_remaining: torch.Tensor | None = None ) -> nemo_automodel.components.speculative.eagle.core_v12.EagleStepMetrics ``` 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. ```python class nemo_automodel.components.speculative.eagle.draft_gpt_oss.GptOssEagle3DraftModel( config: transformers.PretrainedConfig ) ``` **Bases:** [LlamaEagle3DraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama#nemo_automodel-components-speculative-eagle-draft_llama-LlamaEagle3DraftModel) EAGLE-3 draft model for gpt-oss targets. Identical to `LlamaEagle3DraftModel` except that the single draft layer's rotary embedding is replaced with `GptOssDraftRotaryEmbedding` so the draft reproduces gpt-oss's YaRN RoPE instead of the YaRN-incapable `LlamaRotaryEmbedding`. ```python class nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel( model: torch.nn.Module, aux_layer_ids: typing.Sequence[int] | None = None, cp_mesh = None ) ``` **Bases:** [Eagle3TargetBackend](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend) Co-located backend that captures three auxiliary hidden states from a causal LM. **`_cp_size`** `= cp_mesh.size() if cp_mesh is not None else 1` --- **`aux_layer_ids`** `= self._validate_aux_layer_ids(candidate_ids)` --- **`model`** `= model.eval()` --- ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._check_captured( captured: dict[int, torch.Tensor] ) -> None ``` ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._default_aux_layer_ids() -> list[int] ``` ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._get_transformer_layers() -> list[torch.nn.Module] ``` 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. ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._num_hidden_layers() -> int ``` 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. ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel._validate_aux_layer_ids( aux_layer_ids: typing.Sequence[int] ) -> list[int] ``` Validate aux-layer selection before any forward hooks are registered. ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.generate_batch( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, position_ids: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None, doc_remaining: torch.Tensor | None = None ) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch ``` 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. ```python nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel.get_input_embeddings() -> torch.nn.Embedding ``` Return the target model input embeddings. ```python class nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel( model: torch.nn.Module ) ``` Thin wrapper that exposes hidden-state supervision from a causal LM. **`model`** `= model.eval()` --- ```python nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.generate_batch( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, position_ids: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None, doc_remaining: torch.Tensor | None = None ) -> nemo_automodel.components.speculative.eagle.target_v12.EagleTargetBatch ``` 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. ```python nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_input_embeddings() -> torch.nn.Embedding ``` Return the target model input embeddings. ```python nemo_automodel.components.speculative.eagle.target_v12.HFEagleTargetModel.get_lm_head() -> torch.nn.Module ``` Return the target model lm\_head. ```python class nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel( model: torch.nn.Module ) ``` Expose image-aware feature supervision from a frozen Hugging Face VLM. **`image_token_id`** `= int(image_token_id)` --- **`language_backbone`** --- **`model`** `= model.eval()` --- ```python nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.generate_batch( loss_mask: torch.Tensor, model_inputs: dict[str, torch.Tensor] ) -> nemo_automodel.components.speculative.eagle.msd_target.MSDTargetBatch ``` 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:** **`loss_mask`** `torch.Tensor` Bool tensor of shape \[batch, sequence], already aligned to next-token logits. VLM collators produce this from `labels != -100`; they emit `labels[:, 1:]` against `input_ids[:, :-1]`, so no further shift is applied here. Its final position is always dropped because the shifted supervision tensors zero-fill there. --- **`model_inputs`** `dict[str, torch.Tensor]` Mapping of VLM forward inputs. It must include `input_ids` of shape \[batch, sequence], with image-token placeholders at positions that receive projected vision features, and `attention_mask` of the same shape, where one denotes a real token and zero denotes padding. It also carries any vision tensors such as `pixel_values` and `image_grid_thw`. --- **Returns:** `MSDTargetBatch` MSDTargetBatch with shifted image-aware embeddings, target hidden ```python nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel.get_lm_head() -> torch.nn.Module ``` Return the frozen VLM language-model head used by the draft loss. ```python class nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel( config: transformers.PretrainedConfig ) ``` **Bases:** [\_PeagleDraftMixin](/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_draft#nemo_automodel-components-speculative-eagle-peagle_draft-_PeagleDraftMixin), `PreTrainedModel` Llama-style dense EAGLE-3 draft model (Llama, Phi-3, Qwen3). State dict keys match SGLang's `LlamaForCausalLMEagle3` so the saved checkpoint can be loaded by SGLang's inference engine without any remapping (SGLang's `load_weights` fuses `q/k/v_proj` into `qkv_proj` and `gate/up_proj` into `gate_up_proj` via its standard `stacked_params_mapping`). The class name is retained for checkpoint-architectures compatibility; the implementation is config-driven and works for any HF dense decoder-only config that exposes `hidden_size`, `num_attention_heads`, `num_key_value_heads`, `attention_bias`, `mlp_bias`, `rope_theta`, and `rms_norm_eps`. A decoupled `head_dim` is read via `getattr(config, "head_dim", ...)` in the attention layer. Scope: * single draft decoder layer * no KV-cache optimization * no speculative runtime integration **`base_model_prefix`** `= 'model'` --- **`draft_vocab_size`** --- **`has_vocab_compression`** `= self.draft_vocab_size < config.vocab_size` --- **`lm_head`** --- **`model`** `= Eagle3LlamaModel(config)` --- **`target_hidden_size`** --- ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.compute_logits( hidden_states: torch.Tensor ) -> torch.Tensor ``` 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. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.copy_embeddings_from_target( target_embedding: torch.nn.Embedding ) -> None ``` 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. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.embed_input_ids( input_ids: torch.Tensor ) -> torch.Tensor ``` Embed input ids with the draft embedding table. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.forward( input_ids: torch.Tensor, projected_hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor | None = None, cache_hidden: list[list[torch.Tensor]] | None = None, seq_lens: torch.Tensor | None = None ) -> torch.Tensor ``` 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`. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.freeze_embeddings() -> None ``` Freeze draft input embeddings. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_disable() -> None ``` Disable activation checkpointing for the P-EAGLE draft layers. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_enable( gradient_checkpointing_kwargs = None ) -> None ``` 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`. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.project_hidden_states( aux_hidden_states: torch.Tensor ) -> torch.Tensor ``` 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. ```python nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.set_vocab_mapping( selected_token_ids: torch.Tensor ) -> None ``` 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). ```python class nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel( config: transformers.PretrainedConfig ) ``` **Bases:** `PreTrainedModel` Llama-style dense draft that predicts next-step hidden states. Works with Llama, Phi-3, and Qwen3 dense configs. The class name is retained for backward compatibility with already-trained checkpoints. **`embed_tokens`** --- **`fc`** --- **`layers`** --- **`main_input_name`** `= 'input_ids'` --- **`norm`** --- ```python nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.copy_embeddings_from_target( target_embeddings: torch.nn.Embedding ) -> None ``` 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. ```python nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.forward( input_ids: torch.Tensor, target_hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None ) -> torch.Tensor ``` 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. ```python nemo_automodel.components.speculative.eagle.draft_llama_v12.LlamaEagleDraftModel.freeze_embeddings() -> None ``` Freeze draft token embeddings. ```python class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum( config: nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig, global_step: int = 0 ) ``` Stateless-by-step two-stage curriculum with checkpointable progress. Call `source_for_step` from a training loop that owns separate text and multimodal dataloaders, then draw the selected dataloader's next batch. `next_source` is a convenience for single-process loops and stores the next global step in its checkpoint state. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.load_state_dict( state_dict: dict[str, typing.Any] ) -> None ``` Restore curriculum progress from a checkpoint state dictionary. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.multimodal_probability( global_step: int ) -> float ``` Return the multimodal batch probability at `global_step`. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.next_source() -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource ``` Return the current source selection and advance one optimizer step. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.phase_for_step( global_step: int ) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase ``` Return the curriculum phase active at `global_step`. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.source_for_step( global_step: int ) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource ``` 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. ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.state_dict() -> dict[str, typing.Any] ``` Return progress required to resume `next_source`. ```python class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig( text_only_steps: int, multimodal_ramp_steps: int, final_multimodal_ratio: float = 1.0, seed: int = 42 ) ``` Dataclass Configuration for the two-stage MSD data curriculum. The first stage consumes text-only instruction data. In the second stage, the probability of selecting a multimodal batch rises linearly to `final_multimodal_ratio`. A final ratio of one makes the post-ramp phase entirely multimodal, matching the curriculum described in the MSD paper. **Parameters:** **`text_only_steps`** `int` Number of initial optimizer steps using only text batches. --- **`multimodal_ramp_steps`** `int` Number of stage-two steps over which the multimodal sampling probability rises to its final value. Set to zero to switch immediately after the text-only stage. --- **`final_multimodal_ratio`** `float` — default: 1.0 Multimodal batch probability after the ramp. --- **`seed`** `int` — default: 42 Seed used for stateless, step-indexed source selection. All distributed ranks therefore select the same source for a global optimizer step without communicating RNG state. --- **`final_multimodal_ratio`** `float = 1.0` --- **`multimodal_ramp_steps`** `int` --- **`seed`** `int = 42` --- **`text_only_steps`** `int` --- ```python nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig.__post_init__() -> None ``` ```python class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase ``` **Bases:** `enum.Enum` High-level phase of the two-stage MSD training schedule. **`MULTIMODAL`** `= 'multimodal'` --- **`MULTIMODAL_RAMP`** `= 'multimodal_ramp'` --- **`TEXT_ONLY`** `= 'text_only'` --- ```python class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource ``` **Bases:** `enum.Enum` Data source selected for one MSD optimizer step. **`MULTIMODAL`** `= 'multimodal'` --- **`TEXT`** `= 'text'` --- ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder( target: nemo_automodel.components.speculative.eagle.msd_target.HFMSDTargetModel, draft_model: torch.nn.Module ) ``` Batch-one reference MSD round with recursive drafting and tree verification. **`generator`** --- **`target_embeddings`** `= self._resolve_target_embeddings(target.model)` --- ```python nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder._resolve_target_embeddings( model: torch.nn.Module ) -> torch.nn.Module ``` staticmethod Find the language token embedding module across common HF VLM layouts. ```python nemo_automodel.components.speculative.eagle.msd_decode.MSDGreedyDecoder.decode_round( model_inputs: dict[str, torch.Tensor], draft_steps: int, top_k: int, beam_width: int ) -> tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal, nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult] ``` Draft and greedily verify one MSD candidate tree for a VLM prompt. ```python class nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule( draft_model: torch.nn.Module, target_lm_head: torch.nn.Module, hidden_loss_weight: float = 1.0, token_loss_weight: float = 0.1 ) ``` **Bases:** `Module` Draft-side training module for multimodal speculative decoding. **`hidden_loss_fn`** `= nn.SmoothL1Loss(reduction='none')` --- ```python nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.compute_logits( hidden_states: torch.Tensor ) -> torch.Tensor ``` Project predicted hidden states through the frozen target language head. **Parameters:** **`hidden_states`** `torch.Tensor` Tensor of shape \[..., hidden], with arbitrary leading dimensions. --- **Returns:** `torch.Tensor` Tensor of shape \[..., vocab], containing target-vocabulary logits. ```python nemo_automodel.components.speculative.eagle.msd.MSDTrainerModule.forward( inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, input_hidden_states: torch.Tensor, target_hidden_states: torch.Tensor, target_logits: torch.Tensor, image_mask: torch.Tensor ) -> nemo_automodel.components.speculative.eagle.msd.MSDStepMetrics ``` Run one multimodal speculative-decoding training step. **Parameters:** **`inputs_embeds`** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], the target VLM input embeddings shifted left by one position. --- **`attention_mask`** `torch.Tensor` Tensor of shape \[batch, sequence], where one denotes a real token and zero denotes padding. --- **`loss_mask`** `torch.Tensor` Bool tensor of shape \[batch, sequence], true where the shifted next-token target is supervised. --- **`input_hidden_states`** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], the target's unshifted final hidden states used as draft features. --- **`target_hidden_states`** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], the target's final hidden states shifted left by one position. --- **`target_logits`** `torch.Tensor` Tensor of shape \[batch, sequence, vocab], the frozen target logits shifted left by one position. --- **`image_mask`** `torch.Tensor` Bool tensor of shape \[batch, sequence], true where `inputs_embeds` contains a projected image token. --- **Returns:** `MSDStepMetrics` MSDStepMetrics containing scalar loss terms, first-token accuracy, and ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator( draft_model: torch.nn.Module, target_lm_head: torch.nn.Module, target_embeddings: torch.nn.Module ) ``` Recursively draft a top-k MSD feature tree from target prefix features. ```python nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeDraftGenerator.propose( shifted_inputs_embeds: torch.Tensor, input_hidden_states: torch.Tensor, attention_mask: torch.Tensor, shifted_image_mask: torch.Tensor, root_token_id: int, draft_steps: int, top_k: int, beam_width: int ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal ``` 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. ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout( attention_mask: torch.Tensor, position_ids: torch.Tensor, retrieve_indices: torch.Tensor ) ``` Dataclass Tree attention and retrieval metadata for an MSD proposal. Node zero represents the target-sampled root token. `attention_mask` is boolean and contains the root plus every ancestor for each node. Serving backends turn its false entries into additive `-inf` attention bias. **`attention_mask`** `Tensor` --- **`position_ids`** `Tensor` --- **`retrieve_indices`** `Tensor` --- ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode( index: int, parent_index: int, token_id: int, depth: int, log_probability: float ) ``` Dataclass One drafted text-token node in an MSD candidate tree. **`depth`** `int` --- **`index`** `int` --- **`log_probability`** `float` --- **`parent_index`** `int` --- **`token_id`** `int` --- ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal( root_token_id: int, nodes: tuple[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode, ...], leaf_indices: tuple[int, ...], layout: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout ) ``` Dataclass A target root token plus recursively drafted candidate-tree nodes. **`layout`** `MSDTreeLayout` --- **`leaf_indices`** `tuple[int, ...]` --- **`nodes`** `tuple[MSDTreeNode, ...]` --- **`root_token_id`** `int` --- ```python nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_path( leaf_index: int ) -> tuple[int, ...] ``` Return the root-to-leaf candidate-token path. ```python nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal.candidate_paths() -> tuple[tuple[int, ...], ...] ``` Return all leaf paths in deterministic tree order. ```python class nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult( accepted_token_ids: tuple[int, ...], bonus_token_id: int, accepted_draft_tokens: int, leaf_index: int | None ) ``` Dataclass Greedy posterior verification result for one candidate tree. **`accepted_draft_tokens`** `int` --- **`accepted_token_ids`** `tuple[int, ...]` --- **`bonus_token_id`** `int` --- **`emitted_token_ids`** `tuple[int, ...]` Return accepted candidate tokens followed by the target bonus token. --- **`leaf_index`** `int | None` --- ```python class nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel() ``` **Bases:** [LlamaEagleDraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12#nemo_automodel-components-speculative-eagle-draft_llama_v12-LlamaEagleDraftModel) EAGLE-1/2 draft model with modality-aware input feature construction. ```python nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.copy_embeddings_from_target( target_embeddings: torch.nn.Module ) -> None ``` Raise: MSD drafts consume target embeddings instead of holding their own. ```python nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.forward( inputs_embeds: torch.Tensor, target_hidden_states: torch.Tensor, attention_mask: torch.Tensor, image_mask: torch.Tensor ) -> torch.Tensor ``` Predict next-position target hidden states from multimodal features. **Parameters:** **`inputs_embeds`** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], containing the target VLM input embeddings shifted one position to the left. Image positions already contain projected vision features. --- **`target_hidden_states`** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], the frozen target language model's unshifted final hidden states. --- **`attention_mask`** `torch.Tensor` Tensor of shape \[batch, sequence], where one denotes a real token and zero denotes padding. --- **`image_mask`** `torch.Tensor` Bool tensor of shape \[batch, sequence], aligned with `inputs_embeds` and true for projected image-token positions. --- **Returns:** `torch.Tensor` Tensor of shape \[batch, sequence, hidden], the draft predictions for ```python nemo_automodel.components.speculative.eagle.msd.MultimodalEagleDraftModel.freeze_embeddings() -> None ``` Raise: MSD drafts have no token embedding table to freeze. ```python class nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule( draft_model: torch.nn.Module, selected_token_ids: torch.Tensor, selected_token_mask: torch.Tensor, num_depths: int, mask_token_id: int, down_sample_ratio: float = 0.7, down_sample_ratio_min: float = 0.2, sequence_partitions: int = 1 ) ``` **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](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** (`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 `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 (`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: `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. **`down_sample_ratio`** `= float(down_sample_ratio)` --- **`down_sample_ratio_min`** `= float(down_sample_ratio_min)` --- **`mask_token_id`** `= int(mask_token_id)` --- **`sequence_partitions`** `= int(sequence_partitions)` --- ```python 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 ``` 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. ```python 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] ``` 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. ```python nemo_automodel.components.speculative.eagle.peagle_trainer.PEagleTrainerModule.build_peagle_plan( loss_mask: torch.Tensor ) -> '_PeaglePlan' ``` 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. ```python 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 ``` 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. ```python class nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel( runner: nemo_automodel.components.speculative.eagle.target_runner.TargetRunner, aux_layer_ids: typing.Sequence[int] | None = None ) ``` **Bases:** [Eagle3TargetBackend](/nemo-automodel/nemo_automodel/components/speculative/eagle/backend#nemo_automodel-components-speculative-eagle-backend-Eagle3TargetBackend) EAGLE-3 target backend that runs the frozen target through a runner. Engine-agnostic: it owns only the supervision *contract* and delegates the actual forward to a `TargetRunner`. Concrete engines subclass this to add their own `from_pretrained` (which builds the runner); everything else is inherited. **Parameters:** **`runner`** `TargetRunner` A loaded runner implementing `TargetRunner`. --- **`aux_layer_ids`** `Sequence[int] | None` — default: None The three decoder layers to capture (low / mid / high). When `None` the shared EAGLE-3 default recipe is used, matching every other backend. --- **`aux_layer_ids`** --- **`model`** `= runner.model` --- ```python nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.close() -> None ``` Release the runner (frees GPU memory / engine handles). ```python nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.generate_batch( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor ) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch ``` 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. ```python nemo_automodel.components.speculative.eagle.target_runner.RunnerEagle3TargetModel.get_input_embeddings() -> types.SimpleNamespace ``` 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`. ```python class nemo_automodel.components.speculative.eagle.target_runner.TargetRunner() ``` Protocol Minimal engine surface `RunnerEagle3TargetModel` depends on. Implemented for real by an inference-engine runner (`SGLangTargetRunner` today, a vLLM runner later) and faked in unit tests, which is why the backend depends on this protocol rather than on any engine directly. **`model`** `'nn.Module'` --- ```python nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.forward_eagle3( input_ids: torch.Tensor, attention_mask: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor] ``` 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). ```python nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.input_embedding_weight() -> torch.Tensor ``` Return the target input-embedding weight `[vocab, hidden]`. ```python nemo_automodel.components.speculative.eagle.target_runner.TargetRunner.set_aux_layers( aux_layer_ids: typing.Sequence[int] ) -> None ``` Tell the underlying model which 3 decoder layers to capture. ```python nemo_automodel.components.speculative.eagle.msd_decode.accept_or_resample( target_logits: torch.Tensor, draft_logits: torch.Tensor, candidate_token_id: int, generator: torch.Generator | None = None, random_value: float | None = None ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticStep ``` Apply the lossless speculative acceptance rule for one candidate token. `generator` must live on the same device as the logits, because both the acceptance draw and the residual resample run there. `random_value` comes from the half-open interval `[0, 1)` that `torch.rand` samples, and the candidate is accepted on a strict `draw < acceptance` so a token the target assigns zero probability can never be accepted. ```python nemo_automodel.components.speculative.eagle.msd_decode.build_msd_tree_layout( nodes: typing.Iterable[nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeNode], leaf_indices: typing.Iterable[int], device: torch.device ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeLayout ``` Build engine-ready tree attention and root-to-leaf retrieval metadata. ```python nemo_automodel.components.speculative.eagle.registry.resolve_eagle1_draft_spec( architectures: list[str] ) -> nemo_automodel.components.speculative.eagle.registry.DraftSpec ``` Resolve the EAGLE-1 / EAGLE-2 draft spec for a target's `config.architectures` field. ```python nemo_automodel.components.speculative.eagle.registry.resolve_eagle3_draft_spec( architectures: list[str] ) -> nemo_automodel.components.speculative.eagle.registry.DraftSpec ``` Resolve the EAGLE-3 draft spec for a target's `config.architectures` field. ```python nemo_automodel.components.speculative.eagle.core.simulated_accept_length( step_prefix_hits: torch.Tensor, step_valid: torch.Tensor ) -> torch.Tensor ``` 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. ```python nemo_automodel.components.speculative.eagle.msd_decode.verify_greedy_tree( proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal, target_logits_by_leaf: typing.Iterable[torch.Tensor] ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult ``` Select the longest target-greedy accepted path and its bonus token. Each logits tensor has one row for every candidate token in the matching leaf path plus one final row for the target bonus token. This is the greedy form of EAGLE/MSD posterior verification. ```python nemo_automodel.components.speculative.eagle.msd_decode.verify_hf_greedy_tree( model: torch.nn.Module, model_inputs: dict[str, torch.Tensor], proposal: nemo_automodel.components.speculative.eagle.msd_decode.MSDTreeProposal ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDVerificationResult ``` Reference target verification for a batch-one Hugging Face VLM. The path-wise implementation is intentionally engine neutral and works for VLMs whose visual tensors are prompt-only. A serving backend can instead consume `proposal.layout` to verify the same tree in one tree-attention target forward. ```python nemo_automodel.components.speculative.eagle.msd_decode.verify_stochastic_chain( candidate_token_ids: typing.Iterable[int], target_logits: torch.Tensor, draft_logits: torch.Tensor, generator: torch.Generator | None = None ) -> nemo_automodel.components.speculative.eagle.msd_decode.MSDStochasticVerificationResult ``` Losslessly verify a linear draft chain and emit the correction token. ```python nemo_automodel.components.speculative.eagle.registry.EAGLE1_DRAFT_REGISTRY: dict[str, DraftSpec] = {arch: (DraftSpec(draft_cls=LlamaEagleDraftModel)) for arch in _DENSE_ARCHITECTU... ``` ```python nemo_automodel.components.speculative.eagle.registry.EAGLE3_DRAFT_REGISTRY: dict[str, DraftSpec] = {arch: (DraftSpec(draft_cls=LlamaEagle3DraftModel)) for arch in _DENSE_ARCHITECT... ``` ## Docs - [nemo_automodel.components.speculative.eagle.backend](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/backend.md) - [nemo_automodel.components.speculative.eagle.core](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/core.md) - [nemo_automodel.components.speculative.eagle.core_v12](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/core_v12.md) - [nemo_automodel.components.speculative.eagle.draft_deepseek](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_deepseek.md) - [nemo_automodel.components.speculative.eagle.draft_gemma](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gemma.md) - [nemo_automodel.components.speculative.eagle.draft_gpt_oss](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_gpt_oss.md) - [nemo_automodel.components.speculative.eagle.draft_kimi_k3](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_kimi_k3.md) - [nemo_automodel.components.speculative.eagle.draft_llama](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama.md) - [nemo_automodel.components.speculative.eagle.draft_llama_v12](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama_v12.md) - [nemo_automodel.components.speculative.eagle.msd](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/msd.md) - [nemo_automodel.components.speculative.eagle.msd_curriculum](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_curriculum.md) - [nemo_automodel.components.speculative.eagle.msd_decode](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_decode.md) - [nemo_automodel.components.speculative.eagle.msd_target](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/msd_target.md) - [nemo_automodel.components.speculative.eagle.peagle_attention](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_attention.md) - [nemo_automodel.components.speculative.eagle.peagle_data](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_data.md) - [nemo_automodel.components.speculative.eagle.peagle_draft](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_draft.md) - [nemo_automodel.components.speculative.eagle.peagle_trainer](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/peagle_trainer.md) - [nemo_automodel.components.speculative.eagle.registry](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/registry.md) - [nemo_automodel.components.speculative.eagle.remote](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote.md) - [nemo_automodel.components.speculative.eagle.remote.client](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote/client.md) - [nemo_automodel.components.speculative.eagle.remote.protocol](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote/protocol.md) - [nemo_automodel.components.speculative.eagle.remote.server](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote/server.md) - [nemo_automodel.components.speculative.eagle.remote.transport](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote/transport.md) - [nemo_automodel.components.speculative.eagle.remote.wire](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/remote/wire.md) - [nemo_automodel.components.speculative.eagle.ring_attention](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/ring_attention.md) - [nemo_automodel.components.speculative.eagle.sglang_runner](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_runner.md) - [nemo_automodel.components.speculative.eagle.sglang_target](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/sglang_target.md) - [nemo_automodel.components.speculative.eagle.target](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/target.md) - [nemo_automodel.components.speculative.eagle.target_runner](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/target_runner.md) - [nemo_automodel.components.speculative.eagle.target_v12](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/target_v12.md) - [nemo_automodel.components.speculative.eagle.ulysses_attention](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/ulysses_attention.md) - [nemo_automodel.components.speculative.eagle.vispec_core](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_core.md) - [nemo_automodel.components.speculative.eagle.vispec_draft](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_draft.md) - [nemo_automodel.components.speculative.eagle.vispec_target](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/vispec_target.md) - [nemo_automodel.components.speculative.eagle.vllm_runner](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_runner.md) - [nemo_automodel.components.speculative.eagle.vllm_target](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/vllm_target.md) - [nemo_automodel.components.speculative.eagle.zigzag_ring_attention](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/eagle/zigzag_ring_attention.md)