> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-datasets-llm-dspark_cache-CachedDSparkDataset) | Read DSpark offline cache shards lazily. |

### Functions

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

### Data

[`CACHE_KEYS`](#nemo_automodel-components-datasets-llm-dspark_cache-CACHE_KEYS)

[`DTYPE_MAP`](#nemo_automodel-components-datasets-llm-dspark_cache-DTYPE_MAP)

[`_CACHE_NAME`](#nemo_automodel-components-datasets-llm-dspark_cache-_CACHE_NAME)

[`_FORMAT_VERSION`](#nemo_automodel-components-datasets-llm-dspark_cache-_FORMAT_VERSION)

[`_IDENTITY_EXEMPT_FIELDS`](#nemo_automodel-components-datasets-llm-dspark_cache-_IDENTITY_EXEMPT_FIELDS)

[`_TARGET_WEIGHTS_NAME`](#nemo_automodel-components-datasets-llm-dspark_cache-_TARGET_WEIGHTS_NAME)

[`__all__`](#nemo_automodel-components-datasets-llm-dspark_cache-__all__)

### API

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

**Bases:** [CachedTensorDataset](/nemo-automodel/nemo_automodel/components/datasets/llm/offline_cache#nemo_automodel-components-datasets-llm-offline_cache-CachedTensorDataset)

Read DSpark offline cache shards lazily.

```python
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.

```python
nemo_automodel.components.datasets.llm.dspark_cache._target_weights_path(
    cache_dir: str
) -> str
```

```python
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.

```python
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.

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

```python
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.

```python
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.

```python
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.

```python
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`.

```python
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.

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

Persist the DSpark cache manifest.

```python
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.

```python
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.

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

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

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

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

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

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

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