> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.offline_cache

Shared on-disk helpers for speculative offline cache readers.

## Module Contents

### Classes

| Name                                                                                               | Description                                                                    |
| -------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`CachedTensorDataset`](#nemo_automodel-components-datasets-llm-offline_cache-CachedTensorDataset) | Lazily read fixed-key safetensors shards by sample index.                      |
| [`_ShardAccumulator`](#nemo_automodel-components-datasets-llm-offline_cache-_ShardAccumulator)     | Buffer per-batch cache tensors and flush them as fixed-size sequential shards. |

### Functions

| Name                                                                                                                         | Description                                                                             |
| ---------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| [`_sequential_subset_loader`](#nemo_automodel-components-datasets-llm-offline_cache-_sequential_subset_loader)               | Return a sequential loader over `dataset[start:end)` mirroring `dataloader`.            |
| [`atomic_write`](#nemo_automodel-components-datasets-llm-offline_cache-atomic_write)                                         | Write through a sibling temporary file, then atomically replace the target path.        |
| [`build_cached_dataloader`](#nemo_automodel-components-datasets-llm-offline_cache-build_cached_dataloader)                   | Build a dataloader over a precomputed offline cache dataset.                            |
| [`collate_cached`](#nemo_automodel-components-datasets-llm-offline_cache-collate_cached)                                     | Stack per-sample cache dicts into a batch.                                              |
| [`dataloader_from_sample`](#nemo_automodel-components-datasets-llm-offline_cache-dataloader_from_sample)                     | Return an equivalent sequential dataloader beginning at `start_sample`.                 |
| [`ensure_supervision_options_match`](#nemo_automodel-components-datasets-llm-offline_cache-ensure_supervision_options_match) | Reject a cache whose stored supervision was shaped by different options than the run's. |
| [`existing_shard_indices`](#nemo_automodel-components-datasets-llm-offline_cache-existing_shard_indices)                     | Return the shard indices already present in `cache_dir`.                                |
| [`load_safetensors`](#nemo_automodel-components-datasets-llm-offline_cache-load_safetensors)                                 | Return `(save_file, safe_open)` or raise a clear dependency error.                      |
| [`manifest_path`](#nemo_automodel-components-datasets-llm-offline_cache-manifest_path)                                       | Return the manifest path inside `cache_dir`.                                            |
| [`partition_cache_shards`](#nemo_automodel-components-datasets-llm-offline_cache-partition_cache_shards)                     | Assign a contiguous, shard-aligned block of samples to `rank`.                          |
| [`read_manifest`](#nemo_automodel-components-datasets-llm-offline_cache-read_manifest)                                       | Load and validate an offline-cache manifest.                                            |
| [`resume_start_sample`](#nemo_automodel-components-datasets-llm-offline_cache-resume_start_sample)                           | Return the sample index where a deterministic sequential resume should start.           |
| [`shard_path`](#nemo_automodel-components-datasets-llm-offline_cache-shard_path)                                             | Return the path of shard `shard_index` inside `cache_dir`.                              |
| [`validate_contiguous_shards`](#nemo_automodel-components-datasets-llm-offline_cache-validate_contiguous_shards)             | Return shard indices, requiring exactly `0..N-1` for the manifest size.                 |
| [`write_cache_shards`](#nemo_automodel-components-datasets-llm-offline_cache-write_cache_shards)                             | Run a precompute loop and write sequential cache shards.                                |
| [`write_cache_shards_distributed`](#nemo_automodel-components-datasets-llm-offline_cache-write_cache_shards_distributed)     | Distributed precompute loop: each rank writes its contiguous shard block.               |
| [`write_manifest`](#nemo_automodel-components-datasets-llm-offline_cache-write_manifest)                                     | Persist an offline-cache manifest atomically.                                           |
| [`write_tensor_shard`](#nemo_automodel-components-datasets-llm-offline_cache-write_tensor_shard)                             | Write one cache shard containing exactly the requested tensor keys.                     |

### Data

[`MANIFEST_NAME`](#nemo_automodel-components-datasets-llm-offline_cache-MANIFEST_NAME)

[`SHARD_RE`](#nemo_automodel-components-datasets-llm-offline_cache-SHARD_RE)

### API

```python
class nemo_automodel.components.datasets.llm.offline_cache.CachedTensorDataset(
    cache_dir: str,
    cache_name: str,
    format_version: int,
    disk_keys: tuple[str, ...]
)
```

**Bases:** `Dataset`

Lazily read fixed-key safetensors shards by sample index.

**`_open_handles`** `dict[int, Any] = {}`

---

**`_shard_indices`**

---

**`manifest`**

---

**`num_samples`** `= int(self.manifest['num_samples'])`

---

**`shard_size`** `= int(self.manifest['shard_size'])`

---

```python
nemo_automodel.components.datasets.llm.offline_cache.CachedTensorDataset.__getitem__(
    index: int
) -> dict[str, torch.Tensor]
```

```python
nemo_automodel.components.datasets.llm.offline_cache.CachedTensorDataset.__len__() -> int
```

```python
nemo_automodel.components.datasets.llm.offline_cache.CachedTensorDataset._handle(
    shard_index: int
)
```

```python
class nemo_automodel.components.datasets.llm.offline_cache._ShardAccumulator(
    output_dir: str,
    shard_size: int,
    start_shard_index: int,
    write_shard_fn: typing.Callable[[str, int, dict[str, torch.Tensor]], str],
    logger: logging.Logger
)
```

Buffer per-batch cache tensors and flush them as fixed-size sequential shards.

Shard indices count up from `start_shard_index`; a trailing partial shard is
written by `finalize`. Shared by the single-process and distributed
precompute loops so both produce the identical on-disk shard layout.

**`_buffered`** `= 0`

---

**`_chunks`** `list[dict[str, Tensor]] = []`

---

```python
nemo_automodel.components.datasets.llm.offline_cache._ShardAccumulator._flush(
    max_samples: int | None = None
) -> None
```

```python
nemo_automodel.components.datasets.llm.offline_cache._ShardAccumulator.add(
    batch_cache: dict[str, torch.Tensor]
) -> None
```

Append one computed batch, flushing whole shards as they fill.

```python
nemo_automodel.components.datasets.llm.offline_cache._ShardAccumulator.finalize() -> None
```

Write the trailing partial shard, if any.

```python
nemo_automodel.components.datasets.llm.offline_cache._sequential_subset_loader(
    dataloader: torch.utils.data.DataLoader,
    start: int,
    end: int
) -> torch.utils.data.DataLoader
```

Return a sequential loader over `dataset[start:end)` mirroring `dataloader`.

Preserves the source loader's worker settings: in particular
`multiprocessing_context` -- the precompute paths load the target onto CUDA
before iterating, so fork-started workers would inherit a live CUDA context and
abort (see `build_eagle3_dataloader`'s forkserver note).

```python
nemo_automodel.components.datasets.llm.offline_cache.atomic_write(
    path: str,
    write_fn: typing.Callable[[str], None]
) -> str
```

Write through a sibling temporary file, then atomically replace the target path.

```python
nemo_automodel.components.datasets.llm.offline_cache.build_cached_dataloader(
    dataset: torch.utils.data.Dataset,
    batch_size: int,
    shuffle: bool,
    collate_fn: typing.Callable[[list[dict[str, torch.Tensor]]], dict[str, torch.Tensor]],
    num_workers: int = 0,
    distributed: bool = False
) -> torch.utils.data.DataLoader
```

Build a dataloader over a precomputed offline cache dataset.

```python
nemo_automodel.components.datasets.llm.offline_cache.collate_cached(
    features: list[dict[str, torch.Tensor]],
    keys: tuple[str, ...]
) -> dict[str, torch.Tensor]
```

Stack per-sample cache dicts into a batch.

```python
nemo_automodel.components.datasets.llm.offline_cache.dataloader_from_sample(
    dataloader: torch.utils.data.DataLoader,
    start_sample: int
) -> torch.utils.data.DataLoader
```

Return an equivalent sequential dataloader beginning at `start_sample`.

```python
nemo_automodel.components.datasets.llm.offline_cache.ensure_supervision_options_match(
    manifest: dict[str, typing.Any],
    expected: dict[str, bool],
    cache_name: str,
    cache_dir: str,
    producer_name: str
) -> None
```

Reject a cache whose stored supervision was shaped by different options than the run's.

The cached trainers stream the precomputed `loss_mask`/`position_mask` as
stored, so dataset options such as `mask_generation_prompt` only take effect
through the producer. `expected` maps each such option to the recipe's value.
A manifest that does not record an option (a cache older than the option, or
than its manifest field) cannot be verified and is rejected too, unless
`read_manifest` filled the field in from its defaults.

```python
nemo_automodel.components.datasets.llm.offline_cache.existing_shard_indices(
    cache_dir: str
) -> set[int]
```

Return the shard indices already present in `cache_dir`.

```python
nemo_automodel.components.datasets.llm.offline_cache.load_safetensors(
    cache_name: str
)
```

Return `(save_file, safe_open)` or raise a clear dependency error.

```python
nemo_automodel.components.datasets.llm.offline_cache.manifest_path(
    cache_dir: str
) -> str
```

Return the manifest path inside `cache_dir`.

```python
nemo_automodel.components.datasets.llm.offline_cache.partition_cache_shards(
    num_samples: int,
    shard_size: int,
    world_size: int,
    rank: int
) -> tuple[int, int, int]
```

Assign a contiguous, shard-aligned block of samples to `rank`.

Splits the `ceil(num_samples / shard_size)` shards into contiguous runs across
`world_size` ranks (the first `total_shards % world_size` ranks take one
extra shard). Because the split is on whole shards, every shard is owned by
exactly one rank, so ranks can write straight into a shared output directory
with global shard indices and no cross-rank overlap. Ranks that receive no
shards (when `total_shards < world_size`) get `num_local_samples == 0`.

Returns `(start_shard_index, start_sample, num_local_samples)`.

```python
nemo_automodel.components.datasets.llm.offline_cache.read_manifest(
    cache_dir: str,
    cache_name: str,
    format_version: int,
    producer_name: str | None = None,
    defaults: dict[str, typing.Any] | None = None
) -> dict[str, typing.Any]
```

Load and validate an offline-cache manifest.

`defaults` holds fields added after the format was first released, with
the value a producer that predates them effectively used; a manifest that
lacks such a field is read as if it recorded that default, so older caches
keep comparing and training under the configuration they were built with.

```python
nemo_automodel.components.datasets.llm.offline_cache.resume_start_sample(
    existing_shards: set[int],
    shard_size: int
) -> int
```

Return the sample index where a deterministic sequential resume should start.

```python
nemo_automodel.components.datasets.llm.offline_cache.shard_path(
    cache_dir: str,
    shard_index: int
) -> str
```

Return the path of shard `shard_index` inside `cache_dir`.

```python
nemo_automodel.components.datasets.llm.offline_cache.validate_contiguous_shards(
    cache_dir: str,
    cache_name: str,
    num_samples: int,
    shard_size: int
) -> list[int]
```

Return shard indices, requiring exactly `0..N-1` for the manifest size.

```python
nemo_automodel.components.datasets.llm.offline_cache.write_cache_shards(
    dataloader: torch.utils.data.DataLoader,
    output_dir: str,
    shard_size: int,
    start_shard_index: int,
    compute_batch: typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]],
    write_shard_fn: typing.Callable[[str, int, dict[str, torch.Tensor]], str],
    logger: logging.Logger
) -> None
```

Run a precompute loop and write sequential cache shards.

```python
nemo_automodel.components.datasets.llm.offline_cache.write_cache_shards_distributed(
    dataloader: torch.utils.data.DataLoader,
    output_dir: str,
    shard_size: int,
    world_size: int,
    rank: int,
    compute_batch: typing.Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]],
    write_shard_fn: typing.Callable[[str, int, dict[str, torch.Tensor]], str],
    logger: logging.Logger,
    sync_max_steps: typing.Callable[[int], int] | None = None
) -> None
```

Distributed precompute loop: each rank writes its contiguous shard block.

The target model is sharded across all ranks (expert-parallel / FSDP), so every
`compute_batch` call is a collective that all ranks must enter the same number
of times. This rank forwards its own contiguous slice of the dataset (a different
slice per rank -- pure data parallelism, which the MoE all-to-all routes tokens
for), and pads its step count up to the global maximum with dummy forwards whose
output is discarded, keeping the collectives in lockstep.

Writes are idempotent (atomic replace), so re-running after a partial failure
safely recomputes and overwrites rather than resuming. `sync_max_steps` reduces
the per-rank step count to the global max across ranks (an all-reduce MAX); it
defaults to the identity for a single-process run.

```python
nemo_automodel.components.datasets.llm.offline_cache.write_manifest(
    cache_dir: str,
    manifest: dict[str, typing.Any],
    format_version: int
) -> str
```

Persist an offline-cache manifest atomically.

```python
nemo_automodel.components.datasets.llm.offline_cache.write_tensor_shard(
    cache_dir: str,
    shard_index: int,
    samples: dict[str, torch.Tensor],
    keys: tuple[str, ...],
    cache_name: str
) -> str
```

Write one cache shard containing exactly the requested tensor keys.

```python
nemo_automodel.components.datasets.llm.offline_cache.MANIFEST_NAME = 'manifest.json'
```

```python
nemo_automodel.components.datasets.llm.offline_cache.SHARD_RE = re.compile('^shard-(\\d{6})\\.safetensors$')
```