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`	Return the directory holding the draft `config.json` (`path` or its `model/consolidated`).
`_serving_config`	Validate a training-side draft config and return the serving `config.json`.
`_serving_state_dict`	Drop the `model.` prefix and check the tensors match the EAGLE-4 draft exactly.
`export_eagle4_draft`	Write the serving `config.json` and `model.safetensors` of an Automodel EAGLE-4 draft.
`main`	CLI entry point; see the module docstring.

Data

_COPIED_CONFIG_KEYS

_LAYER_KEYS

_MODEL_PREFIX

logger

API

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).

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.

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).

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.

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

CLI entry point; see the module docstring.

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

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

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

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