> 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.datasets.llm.eagle3_cache

On-disk format + reader for the EAGLE-3 offline target-output cache.

This is the SpecForge "offline" training data path: the frozen target model's
per-token supervision (auxiliary hidden states + the draft-vocab target
distribution) is precomputed once and stored on disk, so draft training reads it
back instead of re-running the (large, frozen) target every step.

It is **extremely disk-intensive** -- on the order of tens of MB per sample for
an 8B target (`aux_hidden_states` is `3 * target_hidden_size` wide), i.e.
multiple TB for a large dataset -- and is largely superseded by online training,
where the target forward is cheap relative to the cache I/O. It is kept for
completeness / reproducibility of the SpecForge offline recipe; prefer the online
path unless you are re-training repeatedly on a fixed, bounded dataset.

This module owns the format (so the producer in
`components/speculative/precompute_eagle3.py` and the training-time reader
agree on one schema):

* `<cache_dir>/manifest.json` -- run config + the `selected_token_ids` used
  to build the draft vocabulary (the recipe reuses these instead of rescanning).
* `<cache_dir>/shard-000000.safetensors` -- one shard holds a contiguous block
  of samples, each field stacked along dim 0:
  `input_ids[n,S]`, `attention_mask[n,S]`, `loss_mask[n,S]` (int64),
  `aux_hidden_states[n,S,3H]` (float), `position_mask[n,S,1]` (bool), and the
  draft-vocab target distribution. The distribution is stored in full as
  `target_probs[n,S,draft_vocab]` (float) by default, or -- when the manifest's
  `target_probs_topk` enables it -- as the top-k factorization
  `target_probs_values[n,S,k]` (float) + `target_probs_indices[n,S,k]` (int32),
  which the reader scatters back to full width. Top-k bounds the dominant field's
  footprint at the cost of the dropped tail mass (a disk/fidelity trade-off, off by
  default; see `precompute_eagle3`).

Each `CachedEagle3Dataset` item is exactly the keyword arguments
`Eagle3TrainerModule.forward` consumes on its precomputed-distribution path
(`target_probs` always yielded at full width).

## Module Contents

### Classes

| Name                                                                                              | Description                                                                |
| ------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`CachedEagle3Dataset`](#nemo_automodel-components-datasets-llm-eagle3_cache-CachedEagle3Dataset) | Reads the EAGLE-3 offline cache; each item is one sample's trainer inputs. |

### Functions

| Name                                                                                                                    | Description                                                                                  |
| ----------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- |
| [`_collate_cached`](#nemo_automodel-components-datasets-llm-eagle3_cache-_collate_cached)                               | Stack per-sample cache dicts into a batch.                                                   |
| [`build_cached_eagle3_dataloader`](#nemo_automodel-components-datasets-llm-eagle3_cache-build_cached_eagle3_dataloader) | Build a dataloader over a precomputed EAGLE-3 cache directory.                               |
| [`compress_target_probs`](#nemo_automodel-components-datasets-llm-eagle3_cache-compress_target_probs)                   | Factor a draft-vocab distribution into its top-`topk` `(values, indices)`.                   |
| [`existing_shard_indices`](#nemo_automodel-components-datasets-llm-eagle3_cache-existing_shard_indices)                 | Return EAGLE-3 shard indices already present in `cache_dir`.                                 |
| [`is_compressed`](#nemo_automodel-components-datasets-llm-eagle3_cache-is_compressed)                                   | True when `target_probs_topk` actually compresses (a positive k below the vocab width).      |
| [`manifest_path`](#nemo_automodel-components-datasets-llm-eagle3_cache-manifest_path)                                   | Return the EAGLE-3 manifest path inside `cache_dir`.                                         |
| [`read_manifest`](#nemo_automodel-components-datasets-llm-eagle3_cache-read_manifest)                                   | Load the cache manifest, raising if it is missing or the wrong format version.               |
| [`read_target_embeddings`](#nemo_automodel-components-datasets-llm-eagle3_cache-read_target_embeddings)                 | Load the target input-embedding table written by `write_target_embeddings`.                  |
| [`reconstruct_target_probs`](#nemo_automodel-components-datasets-llm-eagle3_cache-reconstruct_target_probs)             | Scatter stored top-k `(values, indices)` back into a full `[..., draft_vocab]` distribution. |
| [`write_manifest`](#nemo_automodel-components-datasets-llm-eagle3_cache-write_manifest)                                 | Persist the cache manifest atomically (`.tmp` + `os.replace`).                               |
| [`write_shard`](#nemo_automodel-components-datasets-llm-eagle3_cache-write_shard)                                       | Write one shard atomically.                                                                  |
| [`write_target_embeddings`](#nemo_automodel-components-datasets-llm-eagle3_cache-write_target_embeddings)               | Persist the target input-embedding table the draft initializes from.                         |

### Data

[`CACHE_KEYS`](#nemo_automodel-components-datasets-llm-eagle3_cache-CACHE_KEYS)

[`DTYPE_MAP`](#nemo_automodel-components-datasets-llm-eagle3_cache-DTYPE_MAP)

[`_CACHE_NAME`](#nemo_automodel-components-datasets-llm-eagle3_cache-_CACHE_NAME)

[`_COMMON_KEYS`](#nemo_automodel-components-datasets-llm-eagle3_cache-_COMMON_KEYS)

[`_EMBEDDINGS_NAME`](#nemo_automodel-components-datasets-llm-eagle3_cache-_EMBEDDINGS_NAME)

[`_FORMAT_VERSION`](#nemo_automodel-components-datasets-llm-eagle3_cache-_FORMAT_VERSION)

[`_LOSSLESS_PROBS_KEYS`](#nemo_automodel-components-datasets-llm-eagle3_cache-_LOSSLESS_PROBS_KEYS)

[`_TARGET_PROBS_INDICES`](#nemo_automodel-components-datasets-llm-eagle3_cache-_TARGET_PROBS_INDICES)

[`_TARGET_PROBS_VALUES`](#nemo_automodel-components-datasets-llm-eagle3_cache-_TARGET_PROBS_VALUES)

[`_TOPK_PROBS_KEYS`](#nemo_automodel-components-datasets-llm-eagle3_cache-_TOPK_PROBS_KEYS)

### API

```python
class nemo_automodel.components.datasets.llm.eagle3_cache.CachedEagle3Dataset(
    cache_dir: str
)
```

**Bases:** [CachedTensorDataset](/nemo-automodel/nemo_automodel/components/datasets/llm/offline_cache#nemo_automodel-components-datasets-llm-offline_cache-CachedTensorDataset)

Reads the EAGLE-3 offline cache; each item is one sample's trainer inputs.

Shards are opened lazily with `safetensors.safe_open` (memory-mapped) and
sliced per sample, so the full cache is never loaded into memory at once.
Handles are reopened per worker after a DataLoader fork.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.CachedEagle3Dataset.__getitem__(
    index: int
) -> dict[str, torch.Tensor]
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._collate_cached(
    features: list[dict[str, torch.Tensor]]
) -> dict[str, torch.Tensor]
```

Stack per-sample cache dicts into a batch.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.build_cached_eagle3_dataloader(
    cache_dir: str,
    batch_size: int,
    shuffle: bool,
    num_workers: int = 0,
    distributed: bool = False
) -> torch.utils.data.DataLoader
```

Build a dataloader over a precomputed EAGLE-3 cache directory.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.compress_target_probs(
    target_probs: torch.Tensor,
    topk: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Factor a draft-vocab distribution into its top-`topk` `(values, indices)`.

`target_probs` is `[..., draft_vocab]`. The kept values are renormalized to
sum to 1 so the reconstructed distribution stays a proper distribution for soft
cross-entropy. The dropped tail mass is a disk/fidelity trade-off that grows with
how diffuse the target is (it is *not* negligible for a real LLM: top-64 of a
Qwen3-8B next-token distribution keeps only \~0.92 of the mass), so callers choose
`topk` deliberately; see `precompute_eagle3`. `topk` is clamped to the vocab
width.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.existing_shard_indices(
    cache_dir: str
) -> set[int]
```

Return EAGLE-3 shard indices already present in `cache_dir`.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.is_compressed(
    target_probs_topk: int | None,
    draft_vocab_size: int
) -> bool
```

True when `target_probs_topk` actually compresses (a positive k below the vocab width).

```python
nemo_automodel.components.datasets.llm.eagle3_cache.manifest_path(
    cache_dir: str
) -> str
```

Return the EAGLE-3 manifest path inside `cache_dir`.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.read_manifest(
    cache_dir: str
) -> dict[str, typing.Any]
```

Load the cache manifest, raising if it is missing or the wrong format version.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.read_target_embeddings(
    cache_dir: str
) -> torch.Tensor
```

Load the target input-embedding table written by `write_target_embeddings`.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.reconstruct_target_probs(
    values: torch.Tensor,
    indices: torch.Tensor,
    draft_vocab_size: int
) -> torch.Tensor
```

Scatter stored top-k `(values, indices)` back into a full `[..., draft_vocab]` distribution.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.write_manifest(
    cache_dir: str,
    manifest: dict[str, typing.Any]
) -> str
```

Persist the cache manifest atomically (`.tmp` + `os.replace`).

```python
nemo_automodel.components.datasets.llm.eagle3_cache.write_shard(
    cache_dir: str,
    shard_index: int,
    samples: dict[str, torch.Tensor]
) -> str
```

Write one shard atomically.

`samples` carries the `_COMMON_KEYS` plus exactly one `target_probs`
representation: the full `target_probs` (lossless) or the top-k pair
`target_probs_values` + `target_probs_indices`.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.write_target_embeddings(
    cache_dir: str,
    weight: torch.Tensor
) -> str
```

Persist the target input-embedding table the draft initializes from.

The offline training path never loads the target model, but the draft's
`embed_tokens` must still be seeded from the target's embeddings (EAGLE-3
concatenates token embeddings with the carried hidden state), so the
producer stores them once alongside the cache.

```python
nemo_automodel.components.datasets.llm.eagle3_cache.CACHE_KEYS = _COMMON_KEYS + _LOSSLESS_PROBS_KEYS
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache.DTYPE_MAP = {'bf16': torch.bfloat16, 'fp16': torch.float16, 'fp32': torch.float32}
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._CACHE_NAME = 'EAGLE-3'
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._COMMON_KEYS = ('input_ids', 'attention_mask', 'loss_mask', 'aux_hidden_states', 'position_mask...
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._EMBEDDINGS_NAME = 'target_embeddings.safetensors'
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._FORMAT_VERSION = 2
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._LOSSLESS_PROBS_KEYS = ('target_probs',)
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._TARGET_PROBS_INDICES = 'target_probs_indices'
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._TARGET_PROBS_VALUES = 'target_probs_values'
```

```python
nemo_automodel.components.datasets.llm.eagle3_cache._TOPK_PROBS_KEYS = (_TARGET_PROBS_VALUES, _TARGET_PROBS_INDICES)
```