> This page is for version 0.4.0 · 26.04.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.components.speculative.eagle.export_eagle4

Export a trained Automodel EAGLE-4 draft in the `LlamaForCausalLMEagle4` serving layout.

The EAGLE-4 loaders (TokenSpeed's `LlamaForCausalLMEagle4` and the EAGLE-4
engine plugins) refuse unknown config keys and checkpoint tensors that reach no
parameter, while an Automodel checkpoint carries the training-side config (the
target's config plus recipe keys) and wraps the backbone as `model.`. The
export writes:

* `config.json`: `architectures=["LlamaForCausalLMEagle4"]`,
  `model_type="llama"`, `num_hidden_layers=1` and only the draft-structure
  keys (sizes, norms, RoPE, sliding window, aux-layer ids, candidate head).
* `model.safetensors`: every tensor with the `model.` prefix dropped
  (`fc.weight`, `layers.0.*`, `norm.weight`, `embed_tokens.weight`) next
  to the top-level `lm_head.weight`, `efficient_lm_head.weight` and
  `efficient_lm_head_down_proj.weight`. `embed_tokens` is the target's
  (the recipe keeps it frozen), which the serving adapters load in its place.

Usage::

python -m nemo_automodel.components.speculative.eagle.export_eagle4 \
\<checkpoint>/model/consolidated --output-dir \<export_dir>

## Module Contents

### Functions

| Name                                                                                                            | Description                                                                                |
| --------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ |
| [`_resolve_checkpoint_dir`](#nemo_automodel-components-speculative-eagle-export_eagle4-_resolve_checkpoint_dir) | Return the directory holding the draft `config.json` (`path` or its `model/consolidated`). |
| [`_serving_config`](#nemo_automodel-components-speculative-eagle-export_eagle4-_serving_config)                 | Validate a training-side draft config and return the serving `config.json`.                |
| [`_serving_state_dict`](#nemo_automodel-components-speculative-eagle-export_eagle4-_serving_state_dict)         | Drop the `model.` prefix and check the tensors match the EAGLE-4 draft exactly.            |
| [`export_eagle4_draft`](#nemo_automodel-components-speculative-eagle-export_eagle4-export_eagle4_draft)         | Write the serving `config.json` and `model.safetensors` of an Automodel EAGLE-4 draft.     |
| [`main`](#nemo_automodel-components-speculative-eagle-export_eagle4-main)                                       | CLI entry point; see the module docstring.                                                 |

### Data

[`_COPIED_CONFIG_KEYS`](#nemo_automodel-components-speculative-eagle-export_eagle4-_COPIED_CONFIG_KEYS)

[`_LAYER_KEYS`](#nemo_automodel-components-speculative-eagle-export_eagle4-_LAYER_KEYS)

[`_MODEL_PREFIX`](#nemo_automodel-components-speculative-eagle-export_eagle4-_MODEL_PREFIX)

[`logger`](#nemo_automodel-components-speculative-eagle-export_eagle4-logger)

### API

```python
nemo_automodel.components.speculative.eagle.export_eagle4._resolve_checkpoint_dir(
    path: pathlib.Path
) -> pathlib.Path
```

Return the directory holding the draft `config.json` (`path` or its `model/consolidated`).

```python
nemo_automodel.components.speculative.eagle.export_eagle4._serving_config(
    config: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

Validate a training-side draft config and return the serving `config.json`.

**Raises:**

* `ValueError`: if the checkpoint is not a complete EAGLE-4 draft or uses a
  structure the EAGLE-4 loaders do not build.

```python
nemo_automodel.components.speculative.eagle.export_eagle4._serving_state_dict(
    state_dict: dict[str, torch.Tensor],
    config: dict[str, typing.Any]
) -> dict[str, torch.Tensor]
```

Drop the `model.` prefix and check the tensors match the EAGLE-4 draft exactly.

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

Mapping from Automodel draft key to tensor; tensor shapes are
those of `LlamaEagle3DraftModel` and are left unchanged.

---

**`config`** `dict[str, Any]`

The training-side draft config.

---

**Returns:** `dict[str, torch.Tensor]`

Mapping from serving key to the same tensors (no copies).

```python
nemo_automodel.components.speculative.eagle.export_eagle4.export_eagle4_draft(
    checkpoint_dir: str | pathlib.Path,
    output_dir: str | pathlib.Path
) -> pathlib.Path
```

Write the serving `config.json` and `model.safetensors` of an Automodel EAGLE-4 draft.

**Parameters:**

**`checkpoint_dir`** `str | Path`

The consolidated checkpoint directory (or the step
directory that contains `model/consolidated`).

---

**`output_dir`** `str | Path`

Where to write the export; must differ from the checkpoint.

---

**Returns:** `Path`

The output directory.

```python
nemo_automodel.components.speculative.eagle.export_eagle4.main(
    argv: list[str] | None = None
) -> None
```

CLI entry point; see the module docstring.

```python
nemo_automodel.components.speculative.eagle.export_eagle4._COPIED_CONFIG_KEYS = ('attention_bias', 'attention_dropout', 'bos_token_id', 'eagle_config', 'efficie...
```

```python
nemo_automodel.components.speculative.eagle.export_eagle4._LAYER_KEYS = ('hidden_norm.weight', 'input_layernorm.weight', 'post_attention_layernorm.weigh...
```

```python
nemo_automodel.components.speculative.eagle.export_eagle4._MODEL_PREFIX = 'model.'
```

```python
nemo_automodel.components.speculative.eagle.export_eagle4.logger = logging.getLogger(__name__)
```