nemo_automodel.components.datasets.llm.offline_cache

Shared on-disk helpers for speculative offline cache readers.

Module Contents

Classes

Name	Description
`CachedTensorDataset`	Lazily read fixed-key safetensors shards by sample index.
`_ShardAccumulator`	Buffer per-batch cache tensors and flush them as fixed-size sequential shards.

Functions

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

Data

MANIFEST_NAME

SHARD_RE

API

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'])

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

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

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

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.

_buffered

= 0

_chunks

list[dict[str, Tensor]] = []

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

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.

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

Write the trailing partial shard, if any.

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

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.

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.

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.

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.

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

Return the shard indices already present in cache_dir.

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

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

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

Return the manifest path inside cache_dir.

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

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.

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.

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.

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.

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.

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.

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.

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.

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

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