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`	Reads the EAGLE-3 offline cache; each item is one sample’s trainer inputs.

Functions

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

Data

_LOSSLESS_PROBS_KEYS

_TARGET_PROBS_INDICES

_TARGET_PROBS_VALUES

_TOPK_PROBS_KEYS

API

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

Bases: 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.

_compressed

draft_vocab_size

= int(manifest['draft_vocab_size'])

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

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.

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.

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.

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.

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

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

Return the EAGLE-3 manifest path inside cache_dir.

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.

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.

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.

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

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.

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.

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

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

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

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

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

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

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

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

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

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