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

```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 :meth:`finalize`. Shared by the single-process and distributed
precompute loops so both produce the identical on-disk shard layout.

```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.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
) -> dict[str, typing.Any]
```

Load and validate an offline-cache manifest.

```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$')
```