nemo_automodel.components.speculative.precompute_eagle3

Precompute the EAGLE-3 offline target-output cache (SpecForge “offline” path).

The frozen target model’s per-token supervision is the same on every epoch and every run, yet the online recipe recomputes it each step. This script runs the target once over a dataset and writes the supervision to disk; training then reads it back via cached_target_path instead of loading or running the target at all.

============================== READ THIS ================================= This is the SpecForge offline training path. It is disk intensive — the draft-vocab `target_probs` (`draft_vocab` wide) and `aux_hidden_states` (`3 * target_hidden_size` wide) together cost on the order of 100+ MB per sample for an 8B target at seq-len 2048, i.e. multiple TB for a large corpus. Modern practice trains online, where the target forward is cheap next to that I/O; prefer the online recipe unless you re-train a draft repeatedly on a fixed, bounded dataset. `--target-probs-topk` can shrink the `target_probs` field by storing only its top-k mass, but that is a real disk/fidelity trade-off (a real LLM distribution has a fat tail), so it is off by default — see the flag’s help.

Only EAGLE-3 is supported. EAGLE-1/2 supervise on the full-vocab target distribution, which the recipe keeps on the online path only.

Typical usage (single device; large MoE targets that need sharding must use the online path instead):

python -m nemo_automodel.components.speculative.precompute_eagle3
—target-model meta-llama/Llama-3.1-8B-Instruct
—input-data Aeala/ShareGPT_Vicuna_unfiltered
—output-dir /data/eagle3_cache/sharegpt_llama31
—seq-length 2048 —draft-vocab-size 8192
—batch-size 4 —shard-size 256 —dtype bf16

Then point the recipe at it: recipe_args.cached_target_path: /data/eagle3_cache/....

Module Contents

Functions

Name	Description
`_build_parser`	-
`_compute_batch_cache`	Turn one target-model batch into the per-sample tensors the trainer caches.
`_ensure_resume_compatible`	Refuse to `--resume` into a cache produced with a different configuration.
`_run`	Load the target, scan the dataset once, and write the sharded cache. Returns an exit code.
`_validate_args`	Reject invalid CLI values before loading any model.
`main`	CLI entry point. Parses `argv` and returns the process exit code.

Data

logger

API

nemo_automodel.components.speculative.precompute_eagle3._build_parser() -> argparse.ArgumentParser

nemo_automodel.components.speculative.precompute_eagle3._compute_batch_cache(
    target_batch,
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor,
    cache_dtype: torch.dtype,
    target_probs_topk: int
) -> dict[str, torch.Tensor]

Turn one target-model batch into the per-sample tensors the trainer caches.

Reuses _compute_target_distribution — the exact function the online trainer calls — so the cached supervision is numerically identical to the live path. The draft-vocab target_probs is stored in full, or — when target_probs_topk compresses (a positive k below the draft vocab) — as its top-k (values, indices) factorization. Float fields are downcast to cache_dtype; everything is moved to CPU for writing.

nemo_automodel.components.speculative.precompute_eagle3._ensure_resume_compatible(
    cache_dir: str,
    manifest: dict[str, typing.Any],
    existing_shards: set[int]
) -> None

Refuse to --resume into a cache produced with a different configuration.

Every manifest field shapes the shard contents or their addressing: the target_probs columns follow selected_token_ids (which moves with the dataset, shuffle seed, and draft vocab size), sample order follows the dataset, and tensor shapes follow seq_length / dtype. Shards from a mismatched run are indistinguishable by shape alone, so without this check a resume after changing e.g. --input-data or --shuffle-seed would keep the old shards and bless them with the new manifest, silently corrupting the supervision that training reads back.

nemo_automodel.components.speculative.precompute_eagle3._run(
    args: argparse.Namespace
) -> int

Load the target, scan the dataset once, and write the sharded cache. Returns an exit code.

nemo_automodel.components.speculative.precompute_eagle3._validate_args(
    args: argparse.Namespace
) -> None

Reject invalid CLI values before loading any model.

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

CLI entry point. Parses argv and returns the process exit code.

nemo_automodel.components.speculative.precompute_eagle3.logger = logging.getLogger(__name__)