nemo_automodel.components.datasets.llm.dspark_cache

On-disk format and reader for DSpark offline target supervision.

The DSpark online recipe runs a frozen target model every step to capture the intermediate target features used by the draft and the final hidden state used for the TV / confidence losses. This module owns the disk format that lets a precompute job write those tensors once and lets training stream them later via recipe_args.cached_target_path without loading the target model.

Module Contents

Classes

Name	Description
`CachedDSparkDataset`	Read DSpark offline cache shards lazily.

Functions

Name	Description
`_collate_cached`	Stack per-sample cache dicts into a batch.
`_target_weights_path`	-
`build_cache_manifest`	Assemble the DSpark cache manifest.
`build_cached_dspark_dataloader`	Build a dataloader over a precomputed DSpark cache directory.
`compute_batch_cache`	Convert one captured DSpark target batch into the on-disk cache tensors.
`ensure_manifest_complete`	Reject a cache whose producer did not finish writing.
`manifest_mismatch_fields`	Return the manifest keys whose values differ, ignoring bookkeeping fields.
`read_manifest`	Load and validate the DSpark cache manifest.
`read_target_weight_modules`	Load target weights and return module-like objects exposing `.weight`.
`tokenizer_chat_template_sha256`	Return a stable identity for the tokenizer’s effective (post-override) chat template.
`write_manifest`	Persist the DSpark cache manifest.
`write_shard`	Write one DSpark cache shard.
`write_target_weights`	Persist target embedding and lm_head weights for target-free draft initialization.

Data

_IDENTITY_EXEMPT_FIELDS

_TARGET_WEIGHTS_NAME

__all__

API

class nemo_automodel.components.datasets.llm.dspark_cache.CachedDSparkDataset(
    cache_dir: str
)

Bases: CachedTensorDataset

Read DSpark offline cache shards lazily.

nemo_automodel.components.datasets.llm.dspark_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.dspark_cache._target_weights_path(
    cache_dir: str
) -> str

nemo_automodel.components.datasets.llm.dspark_cache.build_cache_manifest(
    target_model: str,
    target_model_type: str,
    target_text_config,
    seq_length: int,
    dtype: str,
    num_samples: int,
    shard_size: int,
    target_layer_ids: list[int],
    train_data_path: str,
    train_split: str | None,
    shuffle_seed: int,
    mask_reasoning_content: bool,
    chat_template_sha256: str
) -> dict[str, typing.Any]

Assemble the DSpark cache manifest.

Single source of the manifest schema: both precompute entry points call this, and train_dspark validates cached training against these fields, so adding or renaming a field here is the one place the contract changes.

Beyond the tensor-shaping target/cache settings, the manifest records the input identity of the run (dataset path/split, shuffle seed, masking settings, effective chat template): a rerun into an existing directory with a different input therefore fails the manifest-compatibility check instead of silently interleaving old and new supervision shard by shard.

nemo_automodel.components.datasets.llm.dspark_cache.build_cached_dspark_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 DSpark cache directory.

nemo_automodel.components.datasets.llm.dspark_cache.compute_batch_cache(
    target_batch,
    cache_dtype: torch.dtype
) -> dict[str, torch.Tensor]

Convert one captured DSpark target batch into the on-disk cache tensors.

Shared by the single-process and distributed precompute entry points so both producers emit byte-identical cache fields (CACHE_KEYS).

nemo_automodel.components.datasets.llm.dspark_cache.ensure_manifest_complete(
    manifest: dict[str, typing.Any],
    cache_dir: str
) -> None

Reject a cache whose producer did not finish writing.

The precompute entry points write the manifest with complete: false before the first shard and flip it to true only after every shard has been written, so an interrupted (or still-running) precompute cannot be consumed as a valid cache. Manifests written before the marker existed have no complete field and are accepted.

nemo_automodel.components.datasets.llm.dspark_cache.manifest_mismatch_fields(
    recorded: dict[str, typing.Any],
    manifest: dict[str, typing.Any]
) -> list[str]

Return the manifest keys whose values differ, ignoring bookkeeping fields.

format_version and the complete marker describe the on-disk state, not the run configuration, so they never count as a mismatch.

nemo_automodel.components.datasets.llm.dspark_cache.read_manifest(
    cache_dir: str,
    allow_incomplete: bool = False
) -> dict[str, typing.Any]

Load and validate the DSpark cache manifest.

allow_incomplete is for the precompute producers themselves (compat checks against a partially written directory); consumers keep the default so an interrupted precompute is never read as a valid cache.

nemo_automodel.components.datasets.llm.dspark_cache.read_target_weight_modules(
    cache_dir: str
) -> tuple[types.SimpleNamespace, types.SimpleNamespace]

Load target weights and return module-like objects exposing .weight.

nemo_automodel.components.datasets.llm.dspark_cache.tokenizer_chat_template_sha256(
    tokenizer
) -> str

Return a stable identity for the tokenizer’s effective (post-override) chat template.

nemo_automodel.components.datasets.llm.dspark_cache.write_manifest(
    cache_dir: str,
    manifest: dict
) -> str

Persist the DSpark cache manifest.

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

Write one DSpark cache shard.

nemo_automodel.components.datasets.llm.dspark_cache.write_target_weights(
    cache_dir: str,
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    dtype: torch.dtype
) -> str

Persist target embedding and lm_head weights for target-free draft initialization.

nemo_automodel.components.datasets.llm.dspark_cache.CACHE_KEYS = ('input_ids', 'loss_mask', 'target_hidden_states', 'target_last_hidden_states')

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

nemo_automodel.components.datasets.llm.dspark_cache._CACHE_NAME = 'DSpark'

nemo_automodel.components.datasets.llm.dspark_cache._FORMAT_VERSION = 1

nemo_automodel.components.datasets.llm.dspark_cache._IDENTITY_EXEMPT_FIELDS = ('format_version', 'complete')

nemo_automodel.components.datasets.llm.dspark_cache._TARGET_WEIGHTS_NAME = 'target_weights.safetensors'

nemo_automodel.components.datasets.llm.dspark_cache.__all__ = ['CACHE_KEYS', 'DTYPE_MAP', 'CachedDSparkDataset', 'build_cache_manifest', 'buil...