llnl / llnl/professor

Add shard-aware distributed batch sampling for multi-sample HDF5 datasets

Open
#7 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2
Forks
1
Avg merge
4d 21h
Merged PRs (30d)
1

Description

The`dataset_type=1` reader supports logical samples stored within multi-sample HDF5 files. Training currently uses the standard shuffled `DistributedSampler` over logical sample indices.

The sampler is valid for distributed shuffling, but it is unaware of which logical samples share an HDF5 file. As a result, consecutive samples assigned to a rank or DataLoader worker may repeatedly transition between shards instead of exploiting the locality provided by the storage layout.

This limits potential reuse of:

- open HDF5 handles;
- HDF5 metadata and chunk-cache state;
- filesystem page-cache contents;
- sequential or batched access within a shard.

#### Proposed change

Add an optional shard-aware batch sampler for manifests containing a shard path and within-shard sample index.

The first implementation should:

- shuffle shard order each epoch;
- shuffle samples within each shard;
- yield a configurable locality window before moving to another shard;
- distribute work across DDP ranks with equal step counts;
- handle partial shards and remainder samples deterministically;
- preserve deterministic `seed + epoch` behavior;
- document padding, duplication, and `drop_last` behavior.

Advanced node-aware ownership can be considered separately after the basic implementation is measured.

#### Potential acceptance criteria

- [ ] Every intended sample is covered according to documented distributed-sampling semantics.
- [ ] Every rank executes the same number of steps.
- [ ] Sampling is deterministic for a fixed seed and epoch.
- [ ] Access-pattern metrics show fewer shard transitions and greater samples-per-shard-visit than the standard sampler.
- [ ] A smoke test confirms comparable training behavior to the standard shuffled sampler.

Contributor guide

No contributing guide indexed for this repository

Research direction

Locate the dataset_type=1 reader, manifest handling for shard paths and within-shard sample indices, and the current DistributedSampler integration. Define and test seed-plus-epoch determinism, equal rank step counts, partial-shard and remainder behavior, padding, duplication, and drop_last semantics. Done includes reduced shard transitions and a smoke test showing comparable training behavior to standard shuffled sampling.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering, distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.