nemo_automodel.components.speculative.precompute_dspark

Precompute the DSpark offline target-supervision cache.

The online DSpark recipe runs a frozen target model every step to capture the intermediate target hidden states consumed by the draft and the final hidden state used by the TV / confidence losses. This script runs that target once and writes those tensors to disk. Training can then set recipe_args.cached_target_path to stream the cache without loading or running the target model.

Typical usage (single device):

python -m nemo_automodel.components.speculative.precompute_dspark
—target-model Qwen/Qwen3-0.6B
—input-data /data/messages.jsonl
—output-dir /data/dspark_cache/qwen3_06b
—seq-length 2048 —batch-size 4 —shard-size 256 —dtype bf16

Module Contents

Functions

Name	Description
`_build_parser`	-
`_ensure_resume_compatible`	Refuse to resume into shards produced with a different configuration.
`_run`	Load the target, scan the dataset once, and write the sharded cache.
`_validate_args`	Reject invalid CLI values before loading the target model.
`main`	CLI entry point. Parses `argv` and returns the process exit code.

Data

_UNSUPPORTED_OFFLINE_TARGET_MODEL_TYPES

logger

API

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

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

Refuse to resume into shards produced with a different configuration.

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

Load the target, scan the dataset once, and write the sharded cache.

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

Reject invalid CLI values before loading the target model.

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

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

nemo_automodel.components.speculative.precompute_dspark._UNSUPPORTED_OFFLINE_TARGET_MODEL_TYPES = (_DEEPSEEK_V4_MODEL_TYPE, _GLM_5_2_MODEL_TYPE, *_MINIMAX_M3_MODEL_TYPES)

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