nemo_automodel.recipes.llm.precompute_dspark_dist

Distributed (multi-node) precompute of the DSpark offline target-supervision cache.

The single-process precompute_dspark script loads the whole target on one box, so very large targets (DeepSeek-V4-Flash, GLM-5.2) that do not fit on a single 8x80GB node cannot be precomputed with it. This entry point loads such a target frozen through the same expert-parallel / FSDP distributed path the training recipe uses (the routed experts shard across ranks), runs it once over the dataset, and writes the same on-disk DSpark cache that train_dspark consumes through recipe_args.cached_target_path — with no live target during draft training.

It is config-driven (it reuses a training-style YAML: the distributed block that shapes the target’s EP/FSDP mesh, plus the recipe_args target / data / draft fields), and is launched with torchrun exactly like multi-node training::

torchrun —nnodes=4 —node-rank=0 —nproc_per_node=8
—master-addr=<NODE0_IP> —master-port=29500
-m nemo_automodel.recipes.llm.precompute_dspark_dist
-c examples/speculative/dspark/deepseek_v4_flash_dspark_precompute.yaml

Each rank forwards a contiguous, shard-aligned slice of the dataset and writes its own global-indexed shards straight into the shared cache_output_dir (the fleet’s shared filesystem), so no post-hoc merge is needed. Small text targets (Qwen3, Gemma4) are also accepted and simply run data-parallel-replicated for throughput. MiniMax M3 (multimodal) is out of scope: the cache schema is text-only.

Module Contents

Functions

Name	Description
`_build_target`	Build the frozen target for capture, dispatching on model type.
`_ensure_output_dir_compatible`	Refuse to write into a directory whose existing manifest describes a different cache.
`_make_sync_max_steps`	Return an all-reduce-MAX reducer over the default group, or the identity.
`_resolve_cache_settings`	Extract and validate the `(output_dir, dtype, shard_size)` cache knobs.
`main`	CLI entry point. Parses `-c <config.yaml>` and runs the precompute.
`run`	Load the (possibly sharded) target and write the distributed DSpark cache.

Data

logger

API

nemo_automodel.recipes.llm.precompute_dspark_dist._build_target(
    cfg,
    recipe_cfg,
    world_size: int,
    device: torch.device,
    compute_dtype: torch.dtype,
    model_type: str,
    target_path: str,
    trust_remote_code: bool
)

Build the frozen target for capture, dispatching on model type.

DeepSeek V4/V4.1, GLM-5.2, and Kimi K3 load through the sharded EP/FSDP path; other single-process text targets (Qwen3, Gemma4) load replicated for data-parallel throughput. Returns (target_config, target_model).

nemo_automodel.recipes.llm.precompute_dspark_dist._ensure_output_dir_compatible(
    output_dir: str,
    manifest: dict[str, typing.Any]
) -> None

Refuse to write into a directory whose existing manifest describes a different cache.

Distributed precompute writes are idempotent (atomic overwrite), so re-running the same config into the same directory safely recomputes. But mixing shards from a different target / dataset / shape silently corrupts the cache, so a mismatched existing manifest is a hard error (use a fresh cache_output_dir). The manifest carries the run’s input identity (dataset path/split, shuffle seed, masking, effective chat template), so a same-shape different-input rerun is also rejected here rather than interleaving old and new supervision. allow_incomplete lets a rerun continue into a directory whose previous run was interrupted.

nemo_automodel.recipes.llm.precompute_dspark_dist._make_sync_max_steps(
    world_size: int,
    device: torch.device
)

Return an all-reduce-MAX reducer over the default group, or the identity.

nemo_automodel.recipes.llm.precompute_dspark_dist._resolve_cache_settings(
    recipe_cfg
) -> tuple[str, str, int]

Extract and validate the (output_dir, dtype, shard_size) cache knobs.

nemo_automodel.recipes.llm.precompute_dspark_dist.main(
    config_path: str | None = None
) -> int

CLI entry point. Parses -c <config.yaml> and runs the precompute.

nemo_automodel.recipes.llm.precompute_dspark_dist.run(
    cfg
) -> int

Load the (possibly sharded) target and write the distributed DSpark cache.

nemo_automodel.recipes.llm.precompute_dspark_dist.logger = logging.getLogger(__name__)