nemo_automodel.recipes.speculative_decoding.train_eagle1

EAGLE-1 / EAGLE-2 training recipe for Llama-style dense LLMs (Llama, Phi-3, Qwen3) and MoE backbones (Qwen3-MoE).

Module Contents

Classes

Name	Description
`TrainEagle1Recipe`	Recipe for EAGLE-1 training on Llama-style dense LLMs (Llama, Phi-3, Qwen3) and MoE backbones (Qwen3-MoE).

Functions

Name	Description
`_build_feature_noise`	Build the EAGLE-1/2 augmentation from the recipe’s fixed-width knob.
`main`	Entrypoint for `TrainEagle1Recipe`.

Data

logger

API

class nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe(
    cfg
)

Bases: SpecDecodeRecipeBase

Recipe for EAGLE-1 training on Llama-style dense LLMs (Llama, Phi-3, Qwen3) and MoE backbones (Qwen3-MoE).

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._compute_metrics(
    batch: dict[str, torch.Tensor]
)

Run the frozen target and the draft over one micro-batch.

Parameters:

batch

dict[str, torch.Tensor]

Dataloader batch whose tensors are already on self.device. Carries input_ids / attention_mask / loss_mask, each of shape [batch, sequence], plus any packing metadata.

Returns:

EagleStepMetrics for this micro-batch. Subclasses that feed the draft

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._finalize_setup(
    recipe_cfg,
    target_path: str,
    wandb_name_prefix: str
) -> None

Build the optimizer, schedule, checkpointer, and logging around a ready trainer module.

Runs once self.trainer_module / self.draft_model / the dataloaders exist, so recipes that assemble those differently (e.g. ViSpec’s VLM target) share the rest of setup() instead of copying it.

Parameters:

recipe_cfg

The recipe_args config node.

target_path

str

Target model id or path, used for checkpoint metadata and the default W&B run name.

wandb_name_prefix

str

Prefix for the default W&B run name.

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._load_extra_state(
    ckpt_dir: str
) -> None

Restore EAGLE-recipe-specific scalars. Subclasses extend this.

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._loss_components(
    metrics
) -> dict[str, float]

Return the per-term losses this recipe logs alongside the total.

Parameters:

metrics

The step metrics returned by _compute_metrics.

Returns: dict[str, float]

Mapping of log-suffix to scalar value, logged as train/<key>.

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._run_eval()

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe._save_extra_state(
    path: str,
    epoch: int
) -> None

Persist EAGLE-recipe-specific scalars. Subclasses extend this.

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe.run_train_validation_loop()

Run the training loop.

nemo_automodel.recipes.speculative_decoding.train_eagle1.TrainEagle1Recipe.setup()

Build target model, draft model, data, optimizer, and trainer module.

nemo_automodel.recipes.speculative_decoding.train_eagle1._build_feature_noise(
    half_width: float
) -> nemo_automodel.components.speculative.eagle.core_v12.FeatureNoiseConfig | None

Build the EAGLE-1/2 augmentation from the recipe’s fixed-width knob.

nemo_automodel.recipes.speculative_decoding.train_eagle1.main(
    config_path: str | None = None
)

Entrypoint for TrainEagle1Recipe.

nemo_automodel.recipes.speculative_decoding.train_eagle1.logger = logging.getLogger(__name__)