Performance tuning for nvalchemi's Zarr-backed Reader, Dataset, and DataLoader pipeline. Use when configuring AtomicDataZarrReader, Dataset, DataLoader, ZarrWriteConfig, or nvalchemi-io-test for training/inference throughput, especially shuffled access, graph-like random access, fused prefetch, pinned memory, validation overhead, or Zarr chunk/shard choices.
Performance tuning for nvalchemi's Zarr-backed Reader, Dataset, and DataLoader pipeline. Use when configuring AtomicDataZarrReader, Dataset, DataLoader, ZarrWriteConfig, or nvalchemi-io-test for training/inference throughput, especially shuffled access, graph-like random access, fused prefetch, pinned memory, validation overhead, or Zarr chunk/shard choices.
Zarr DataLoader Performance Tuning
Use this skill when optimizing nvalchemi Zarr reads or writing stores that will
later be read through the nvalchemi DataLoader.
Overview
The pipeline has clean ownership boundaries:
Reader: storage I/O only. Returns raw CPU tensor dictionaries plus metadata.
Dataset: validation, optional validation skipping, device transfer, and async
prefetch orchestration. Its canonical explicit batch API is
load_batches(batch_index_lists).
DataLoader: sampler/batch iteration, fused prefetch, stream usage, and batch
construction.
MultiDataset: global index composition over multiple Datasets while routing
load_batches requests to child datasets.
Sampler / batch_sampler: semantic sample order and batch membership. Do not
rely on sampler windows to optimize storage I/O.
Reader public methods:
reader.read(index): one sample.
reader.read_many(indices): many samples, returned in the request order.
Reader backend hooks:
_load_sample(index): implement for simple single-sample formats.
_load_many_samples(indices): implement for batch-optimized formats.
__len__(): total logical samples.
The base Reader owns metadata finalization and optional pinned memory. Index
validity is the concrete reader's responsibility. AtomicDataZarrReader supports
negative logical indices, maps through the active sample mask, and implements
_load_many_samples as the fast path.
Recommended DataLoader setup
from nvalchemi.data.datapipes import (
AtomicDataZarrReader,
Dataset,
DataLoader,
)
reader = AtomicDataZarrReader("store.zarr")
dataset = Dataset(
reader,
device="cuda",
num_workers=1, # 1 is enough; concurrent Zarr reads contend
skip_validation=True, # safe when store was written by the toolkit
)
loader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
prefetch_factor=16, # up to 64 * 16 = 1024 indices per backend read
num_streams=2,
use_streams=True,
pin_memory=True, # request pinned CPU tensors from the reader
)
Use pin_memory=True on AtomicDataZarrReader(...) directly only for manual
reader usage. For normal training, prefer DataLoader(..., pin_memory=True) so
the loader owns the transfer optimization.
Key knobs
prefetch_factor (DataLoader)
Controls how many emitted batches are fused into one backend read:
For batch_size=64, prefetch_factor=16, the model still receives batches of 64
graphs, but the Zarr reader sees up to 1024 logical indices per read_many.
Access pattern
Recommended prefetch_factor
Sequential
2-4
Shuffled
16-64
Block-shuffle
2-8
Use prefetch_factor=0 to disable fused prefetch and issue one backend read per
emitted batch through Dataset.load_batches([indices]). This is useful for
debugging or for stores where larger windows do not help. Positive
prefetch_factor values use the async
prefetch_fused_batches(...) / get_fused_batches() path.
Bypasses per-sample AtomicData Pydantic validation (~4 ms/sample).
Constructs Batch directly from raw tensor dicts via
Batch.from_raw_dicts().
Use when: the store was written by AtomicDataZarrWriter or has been
validated externally.
Do not use when: the store contents are untrusted or from a third party.
num_workers (Dataset)
Thread pool size for background Dataset prefetch work. Start with 1.
Increase only if profiling shows CPU-side validation or device transfer is
underlapping and storage reads are not contending.
pin_memory (DataLoader or Reader)
Pinned CPU tensors make async CPU-to-GPU transfer possible. Use with CUDA targets
and use_streams=True.
chunk_size is rows along dimension 0, not number of structures. Atom fields
are stored on the total atom axis; edge fields on the total edge axis.
Smaller chunks reduce single-sample read amplification but increase metadata
and codec overhead.
Sharding groups many chunks into fewer storage objects and is useful when small
chunks would create too many files.
Use edge_chunk_size / edge_shard_size in nvalchemi-io-test when edge
arrays need different tuning from atom/system arrays.
Zstd level 3 is a good default ratio/speed tradeoff. LZ4 is useful when write
and decompression speed matter more than compression ratio.
How the reader optimises random access
AtomicDataZarrReader._load_many_samples(indices) is the optimized path behind
public reader.read_many(indices).
It currently:
Resolves logical indices through the active sample mask.
Sorts requests by physical sample index.
Groups physical positions by Zarr chunk locality.
Uses coalesced range reads when a small number of chunk-local runs exists.
Falls back to orthogonal selection for highly fragmented requests.
Restores the caller's original request order.
This is transparent to Dataset, DataLoader, and Samplers. Larger fused read
windows give the Zarr backend more indices to coalesce, which is why
prefetch_factor matters most for shuffled reads.
For multidataset training, use MultiDatasetBatchSampler or
MultiDatasetBatchSampler.balanced(...) to define semantic dataset mixing
rates.
samples_per_dataset may be integer counts or float ratios. Use
epoch_policy="max_size", replacement=True when smaller datasets should be
oversampled so the largest dataset does not dominate an epoch.
Benchmark workflow
Use the current CLI subcommands:
# Self-contained write + read benchmark.env COLUMNS=240 uv run nvalchemi-io-test roundtrip \
-n 10000 \
--read-mode batch \
--read-order shuffle \
--batch-size 64 \
--prefetch-factor 16 \
--pin-memory
# Sweep prefetch factors on the same access pattern.for pf in 8 16 32 64 128; doenv COLUMNS=240 uv run nvalchemi-io-test roundtrip \
-n 10000 \
--read-mode batch \
--read-order shuffle \
--batch-size 64 \
--prefetch-factor "$pf" \
--pin-memory
done# Benchmark an existing store without rewriting it.env COLUMNS=240 uv run nvalchemi-io-test read /path/to/store.zarr \
--read-order shuffle \
--batch-size 64 \
--prefetch-factor 32 \
--pin-memory
# Compare DataLoader fused reads against one-sample-at-a-time reads.env COLUMNS=240 uv run nvalchemi-io-test read /path/to/store.zarr \
--read-mode both \
--read-order shuffle \
--batch-size 64 \
--prefetch-factor 32
Important benchmark semantics:
read-mode=batch uses the public DataLoader path with fused prefetch.
Benchmark batch mode uses Dataset(skip_validation=True) to focus on storage
and batching throughput.
read-mode=single calls reader.read(index) once per sample and is only a
baseline for one-sample-at-a-time access.
batch_size is the model-facing batch size.
prefetch_factor controls the backend read window.
Use read-order=shuffle to model fully shuffled training reads.
Use read-order=block-shuffle to test partial locality.
Diagnosing bottlenecks
Run nvalchemi-io-test read on an existing representative store.
Sweep prefetch_factor at the target batch_size.
Compare read-mode=batch against read-mode=single.
If batch mode is fast but training is slow, inspect validation, batching, and
device-transfer overhead. Try skip_validation=True, pin_memory=True, and
CUDA streams.
If batch mode is slow, inspect chunk/shard configuration, compression codec,
filesystem metadata pressure, and read order.
Quick checklist
Use Dataset(skip_validation=True) for trusted toolkit-written stores.
Use DataLoader(pin_memory=True) for CUDA training.
Start with batch_size=64.
Start with prefetch_factor=16 or 32 for shuffled reads.
Sweep prefetch_factor=8,16,32,64,128 with nvalchemi-io-test.
Keep sampler semantics independent from storage locality.
Use load_batches(...) for explicit batch reads.
Tune chunk/shard sizes on a representative store and filesystem.
Use read-mode=single only as a baseline, not as the training path.
