Domain-parallel (rank-local) reading in the datapipes - #1982
coreyjadams wants to merge 1 commit into
Conversation
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 16 file(s)
⏳ @negin513 — 15 file(s)
No CODEOWNER
Comment |
 The PR is not yet safe to merge because Zarr single-group datasets containing scalar-per-sample fields fail before loading. Findings
Summary
Reviews (1) · Last reviewed commit: "Domain-parallel (rank-local) reading in ..." |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| n_rows = ( | ||
| len(window) | ||
| if window is not None and field in windowed_keys | ||
| else self._array_rows(array) | ||
| ) |
There was a problem hiding this comment.
In single-group mode, a scalar field for each sample is stored as an array shaped (n_samples,). This code now unconditionally reads array.shape[1] while building the selection plan, so these fields raise IndexError before any sample is loaded. Previously, they were read directly with root[field][index]. Arrays without a second dimension need to remain supported as replicated scalar fields.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## dp-03-routed-gather #1982 +/- ##
=======================================================
+ Coverage 71.88% 71.94% +0.05%
=======================================================
Files 929 930 +1
Lines 70111 70544 +433
Branches 10560 10632 +72
=======================================================
+ Hits 50400 50752 +352
- Misses 16250 16324 +74
- Partials 3461 3468 +7
🚀 New features to boost your workflow:
|
a5aa9c5 to
5194955
Compare
5194955 to
ee97337
Compare
Readers can produce sharded samples directly: each rank reads only its share of every sharded batch axis from disk, the sample travels as a ShardedProto, and the dataset assembles Shard(0) ShardTensors on the device. Nothing the size of the full sample exists on any rank. - _domain_parallel.py: DomainParallelConfig parsed once from the `domain_parallel` dict + 1-D device_mesh (auto_shard_size gate on the length of dim 0, `placements` overrides by dotted prefix, unmatched keys warn); ShardedProto payload and the communication-free ShardTensor wrap. - Reader: `_load_sample_domain_parallel` contract for custom readers, `_selection_plan` shared by the plain and rank-local reads; ZarrReader and TensorStoreZarrReader implement it with one read procedure each. DomainParallelConfig and resolve_leaf_placements are exported from physicsnemo.datapipes.readers for custom readers. - MeshReader / DomainMeshReader: one batch axis per mesh (cells, with vertices following; points for a point cloud), even splits of the cell and point axes, cells keep global vertex ids so points[cells] is the routed gather; cell windows compact globally as the eager reader does. Memmap and zarr share `_read_mesh_selection`; extra boundaries pin to replicate; subsampling requires a seed. Shared reader base for seed, epoch, zarr group cache and close(). - CenterMesh materializes the center of mass on sharded meshes. - Tests: CPU config/plan tests; GPU datapipe tests for zarr, memmap and zarr meshes, domain meshes (nested global_data, extra boundaries, drop flags, placement overrides, subsample equivalence). Minimal example. User-guide docs follow in a separate PR.
ee97337 to
12ae16b
Compare
| domain_parallel : dict[str, Any], optional | ||
| Optional dict to configure domain-parallel (rank-local) reading. | ||
| Requires ``device_mesh``. If provided, may contain: | ||
|
|
||
| - ``auto_shard_size``: auto gate; a batch axis shards when its | ||
| length (the size of tensor dim 0) is at least this many | ||
| entries (default 1024), e.g. ``(200_000, 3)`` points shard | ||
| while a ``(512, 512)`` image replicates | ||
| - ``placements``: ``{key: "shard" | "replicate"}`` overrides; a | ||
| key pins the whole batch axis (every key sharing its dim-0 | ||
| length), applies by prefix to nested keys, and unnamed axes | ||
| fall back to the gate | ||
|
|
||
| Supporting readers read only this rank's chunk of every sharded | ||
| key and return a proto payload the dataset assembles into | ||
| ``Shard(0)`` ShardTensors on the GPU. Composes with | ||
| ``coordinated_subsampling``: the rank chunk is taken of the | ||
| coordinated window, which requires a seed (``set_generator``) so | ||
| every rank draws the same window. Every rank in the device mesh | ||
| must request the same sample indices: use a sampler that hands | ||
| the same index sequence to every rank of a domain group (e.g. a | ||
| ``DistributedSampler`` over the data-parallel axis only). Readers | ||
| that don't support domain-parallel reading raise. | ||
| device_mesh : torch.distributed.device_mesh.DeviceMesh, optional |
There was a problem hiding this comment.
This is defining the way we let users configure, in yaml, the domain parallel components: they can set automatic parallelism and a threshold to decide how small to replicate; there is also a way to set key-based choices for parallelism
| for key, shape in shapes.items(): | ||
| windowed = window is not None and key in windowed_keys | ||
| if config is not None and shard[key]: | ||
| lo, hi = config.chunk_bounds(shape[0]) | ||
| # This rank's share of the window (a sub-slice of a 1-2 run | ||
| # cyclic block) or of the full range. | ||
| selection[key] = window[lo:hi] if windowed else slice(lo, hi) | ||
| sharded[key] = shape | ||
| else: | ||
| selection[key] = window if windowed else slice(None) | ||
| return selection, sharded |
There was a problem hiding this comment.
When reading selections, in domain parallelism we have to read only a piece of a selection. This function takes a whole region lo:hi and chunks it, maps by domain rank, and spits out "here's the slice selection for this rank"
| Any | ||
| The assembled sample, or *data* unchanged. | ||
| """ | ||
| return assemble_if_proto(data) |
There was a problem hiding this comment.
This is new for domain paralllelism: put it together as a ShardTensor if it's domain parallel.
| r"""Shared machinery for domain-parallel (rank-local) reading in datapipes. | ||
|
|
||
| Readers that support domain-parallel reading accept a ``domain_parallel`` | ||
| configuration dict plus a 1-D ``device_mesh`` (constructed and injected at | ||
| runtime in Python -- it is not serializable config). Each rank reads only | ||
| its share of the sharded batch axes; the local pieces move to the GPU inside a | ||
| :class:`ShardedProto`; and :meth:`ShardedProto.assemble` builds the sample | ||
| with ``Shard(0)`` ShardTensors via the communication-free chunk path of | ||
| ``ShardTensor.from_local``. |
There was a problem hiding this comment.
This file exists to carry two pieces for domain parallel datapipes:
- All the configuration, and it's helper machinery, is consolidated here.
- a reusable "proto" class is here.
The reason we have a proto class is because of the design of ShardTensor vs. datapipes: ShardTensor is GPU native, by choice, but the datapipes draw a fairly firm line between "reading data is CPU" and "preprocessing data is GPU". So to do a domain-parallel read on CPU we can not construct a ShardTensor off the bat: we have to make something else, move it to GPU, and then create ShardTensor objects. That's the ShardProto object. We reuse the same proto object for every datapipe, though, it's quite flexible.
The assemble function is what turns it from the proto -> ShardTensor proper.
|
|
||
| def _load_sample(self, index: int) -> dict[str, torch.Tensor]: | ||
| """Load a single sample from a Zarr group using TensorStore.""" | ||
| def _read_sample( |
There was a problem hiding this comment.
I don't know why claude felt the need to rename this ugh. Going to fix.
| return _read(src.cells[selection]) | ||
|
|
||
|
|
||
| def _read_mesh_selection( |
There was a problem hiding this comment.
This is a wrapper to make one entry point to reading cells, or points, with zarr or memmap.
1 participant
PhysicsNeMo Pull Request
This PR enables domain parallelism as a first-class citizen in many datapipes. It's a big PR, with many pieces, and supports most but not all of the readers.
One major piece of this is the concept of separate CPU and GPU payloads: the datapipes load to CPU, first, then stage to GPU. ShardTensor is GPU only,
so we have a "proto" class here that understands it is a CPU payload that will
become a ShardTensor on the GPU.
By default, Shard(0) only is support.
_domain_parallel.py:DomainParallelConfigparsed once from thedomain_paralleldict plus a 1-Ddevice_mesh(auto_shard_sizegate onthe length of dim 0,
placementsoverrides by dotted prefix, unmatchedkeys warn);
ShardedProtopayload and the communication-free ShardTensorwrap.
Reader:_load_sample_domain_parallelcontract for custom readers and a_selection_planshared by the plain and rank-local reads.ZarrReaderand
TensorStoreZarrReaderimplement it with one read procedure each.DomainParallelConfigandresolve_leaf_placementsare exported forcustom readers.
MeshReader/DomainMeshReader: one batch axis per mesh (cells, withvertices following; points for a point cloud), even splits, cells keep
global vertex ids so
points[cells]is the routed gather, cell windowscompact globally as the eager reader does. Memmap and zarr share
_read_mesh_selection; extra boundaries pin to replicate; subsamplingrequires a seed so every rank of a domain group draws the same window.
CenterMeshmaterializes the center of mass on sharded meshes.meshes and domain meshes (nested
global_data, extra boundaries, dropflags, placement overrides, subsample equivalence). Minimal example under
examples/minimal/datapipes/.Description
Checklist
Dependencies
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.
Stack created with GitHub Stacks CLI • Give Feedback 💬