> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/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.recipes.llm.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`](#nemo_automodel-recipes-llm-train_eagle1-TrainEagle1Recipe) | Recipe for EAGLE-1 training on Llama-style dense LLMs (Llama, Phi-3, Qwen3) and MoE backbones (Qwen3-MoE). |

### Functions

| Name                                                                                          | Description                                                                           |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_all_reduce_mean`](#nemo_automodel-recipes-llm-train_eagle1-_all_reduce_mean)               | -                                                                                     |
| [`_build_feature_noise`](#nemo_automodel-recipes-llm-train_eagle1-_build_feature_noise)       | Build the EAGLE-1/2 augmentation from the recipe's fixed-width knob.                  |
| [`_packing_kwargs`](#nemo_automodel-recipes-llm-train_eagle1-_packing_kwargs)                 | Sequence-packing metadata from a dataloader batch (empty dict when unpacked).         |
| [`_submesh_or_none`](#nemo_automodel-recipes-llm-train_eagle1-_submesh_or_none)               | Return the named (flattened) submesh, or None if absent / no mesh.                    |
| [`_validate_packing_gates`](#nemo_automodel-recipes-llm-train_eagle1-_validate_packing_gates) | Reject sequence-packing configs the EAGLE-1/2 path cannot honor (fail fast at setup). |
| [`main`](#nemo_automodel-recipes-llm-train_eagle1-main)                                       | Entrypoint for `TrainEagle1Recipe`.                                                   |

### Data

[`logger`](#nemo_automodel-recipes-llm-train_eagle1-logger)

### API

```python
class nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe(
    cfg
)
```

**Bases:** [BaseRecipe](/nemo-automodel/nemo_automodel/recipes/base_recipe#nemo_automodel-recipes-base_recipe-BaseRecipe)

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

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._build_checkpointer(
    target_path: str
) -> None
```

Build the checkpointer using the same plumbing as the standard recipes.

```python
nemo_automodel.recipes.llm.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

```python
nemo_automodel.recipes.llm.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.

---

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._load_extra_state(
    ckpt_dir: str
) -> None
```

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

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._log_saved_checkpoint(
    kind: str,
    epoch: int,
    step: int
) -> None
```

Log a saved checkpoint on rank 0 when checkpointing is enabled.

```python
nemo_automodel.recipes.llm.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>`.

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._maybe_save_final_checkpoint(
    completed_epochs: int
) -> bool
```

Always save the fully-trained model at the end of a completed run,
unless a periodic checkpoint already captured the final step.

The end-of-run state is otherwise easy to lose: with no cadence nothing is
saved at all, and with a pure step cadence the final step is skipped
whenever the total step count is not a multiple of `ckpt_every_steps`.
This is a no-op only when a step or epoch checkpoint already landed on the
final step, so it never duplicates or collides with one.

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._maybe_save_step_checkpoint(
    epoch: int
) -> bool
```

Save a checkpoint mid-epoch when `ckpt_every_steps` is configured.

Called after every optimizer step. Saves whenever `ckpt_every_steps` is
a positive integer and the current `global_step` is a multiple of it.
Returns True if a checkpoint was written. The checkpoint directory is
named `epoch_{epoch}_step_{global_step}` so it never collides with the
end-of-epoch checkpoint (which uses `epoch + 1`).

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._module()
```

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._run_eval()
```

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

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

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe._wandb_log(
    data: dict,
    step: int
) -> None
```

Log a metrics dict to W\&B when a run is active (rank 0).

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe.load_checkpoint(
    restore_from: str | None = None
) -> None
```

Resolve and restore a checkpoint produced by `save_checkpoint`.

Restores the draft model, optimizer, LR scheduler, RNG, and `global_step`.
Target model weights are NOT restored — they are re-loaded from the HF hub on
each run because the target is frozen.

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe.run_train_validation_loop()
```

Run the training loop.

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe.save_checkpoint(
    epoch: int,
    step: int,
    train_loss: float | None = None,
    val_loss: dict[str, float] | None = None,
    best_metric_key: str = 'default',
    is_final_checkpoint: bool = False
) -> None
```

Persist draft model, optimizer, scheduler, RNG, and EAGLE meta.

Overrides `BaseRecipe.save_checkpoint` because EAGLE recipes hold multiple
`nn.Module` attributes (frozen target, target wrapper, trainer module wrapping
the draft) — only `draft_model` should be persisted as the main model.

`is_final_checkpoint` is computed by the caller (this hand-rolled loop
has no `step_scheduler` for the checkpointer to infer it from);
`save_consolidated: final` exports HF safetensors only when it is True.

```python
nemo_automodel.recipes.llm.train_eagle1.TrainEagle1Recipe.setup()
```

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

```python
nemo_automodel.recipes.llm.train_eagle1._all_reduce_mean(
    value: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.recipes.llm.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.

```python
nemo_automodel.recipes.llm.train_eagle1._packing_kwargs(
    batch: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]
```

Sequence-packing metadata from a dataloader batch (empty dict when unpacked).

The packed loader (`packed_sequence_size > 0`) emits `position_ids` /
`seq_lens` / `doc_remaining` alongside `input_ids`; the default loader
does not. Keyed on `seq_lens` so the caller can splat the result into
`HFEagleTargetModel.generate_batch` unconditionally.

```python
nemo_automodel.recipes.llm.train_eagle1._submesh_or_none(
    device_mesh,
    name: str
)
```

Return the named (flattened) submesh, or None if absent / no mesh.

Uses `get_flat_mesh` so `_flatten()`-created axes ("dp") resolve across
torch versions. The "dp" axis excludes "tp", so keying the draft DDP group
and the dataloader sampler on it replicates the draft across tensor-parallel
ranks (every TP rank in a draft replica sees the same batch).

```python
nemo_automodel.recipes.llm.train_eagle1._validate_packing_gates(
    cp_size: int,
    target_attn_impl: str,
    micro_batch_size: int
) -> None
```

Reject sequence-packing configs the EAGLE-1/2 path cannot honor (fail fast at setup).

* Context parallelism shards the sequence and strips the 4D block-causal mask
  packing relies on, and EAGLE-1/2 has no CP sequence-sharding path, so
  `cp_size > 1` with packing would silently train on wrong supervision.
* A FlashAttention target infers document boundaries from per-document
  `position_ids`, which transformers packs only at batch size 1.

```python
nemo_automodel.recipes.llm.train_eagle1.main(
    config_path: str | None = None
)
```

Entrypoint for `TrainEagle1Recipe`.

```python
nemo_automodel.recipes.llm.train_eagle1.logger = logging.getLogger(__name__)
```