Standalone datapipe and ShardTensor fixes ahead of domain-parallel reading - #1979
coreyjadams wants to merge 2 commits into
Conversation
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 8 file(s)
⏳ @negin513 — 8 file(s)
No CODEOWNER
Comment |
 The PR is not yet safe to merge because shard-shape inference fails for CPU/Gloo meshes on CUDA-capable hosts. Findings
Summary
Reviews (1) · Last reviewed commit: "Standalone datapipe and ShardTensor fixe..." |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Collectives run on the process group's backend device: CUDA when | ||
| # available, otherwise CPU (gloo) -- so CPU-only single-process use of | ||
| # ``sharding_shapes="infer"`` works too. | ||
| device = "cuda" if torch.cuda.is_available() else "cpu" |
There was a problem hiding this comment.
On a CUDA-capable host using a CPU DeviceMesh backed by Gloo, this selects CUDA for the shape tensors without checking local_group. Calling ShardTensor.from_local(..., sharding_shapes="infer") then passes CUDA tensors to a CPU process group, causing shard-shape inference to fail. Select the collective device from the mesh or process-group backend instead of host-wide CUDA availability.
be1c1eb to
68e8de3
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1979 +/- ##
==========================================
- Coverage 72.64% 72.35% -0.30%
==========================================
Files 928 928
Lines 69963 70005 +42
Branches 10542 10549 +7
==========================================
- Hits 50828 50650 -178
- Misses 15632 15868 +236
+ Partials 3503 3487 -16
🚀 New features to boost your workflow:
|
68e8de3 to
44fc861
Compare
`coverage run -m pytest --testmon` measured nothing executed inside tests: pytest-testmon installs its own coverage.py tracer per test to track dependencies, which replaces the outer tracer. The merged PR report was therefore the nightly baseline (main's line numbers) plus import-time lines, so any file whose lines shifted appeared to lose coverage (e.g. PR #1979: mesh.py lines hit by 28 passing tests reported missing). Reproduced locally on test_mesh_readers.py: plain run 60%, --testmon 15%, --testmon --testmon-nocollect 60%. Pass --testmon-nocollect so testmon only selects tests. The Coverage job restores but never publishes the testmon DB, so nothing is lost. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
`coverage run -m pytest --testmon` measured nothing executed inside tests: pytest-testmon installs its own coverage.py tracer per test to track dependencies, which replaces the outer tracer. The merged PR report was therefore the nightly baseline (main's line numbers) plus import-time lines, so any file whose lines shifted appeared to lose coverage (e.g. PR #1979: mesh.py lines hit by 28 passing tests reported missing). Reproduced locally on test_mesh_readers.py: plain run 60%, --testmon 15%, --testmon --testmon-nocollect 60%. Pass --testmon-nocollect so testmon only selects tests. The Coverage job restores but never publishes the testmon DB, so nothing is lost. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
…ading - MeshReader/DomainMeshReader.close(); extra-boundary loader split - Zarr readers: extract-method refactor, shared Reader._window_indices, deterministic subsampling key order - record_consumer_stream unwraps ShardTensor/DTensor leaves - ShardTensor: scatter_tensor owns its storage, immutable redistribute planner key, spec hash invalidation, gloo shape gather, compile fixes
Datapipes (CPU-runnable): - Reader._window_indices: disabled, list-order key fallback, no target present, seeded reproducibility; include_index_in_metadata toggle. - ZarrReader._array_rows in both modes; absent leading target key falls back to the next configured key; no-target loads full arrays. - TensorStoreZarrReader: same fallbacks plus _finalize_sample merging arrays, attributes and defaults. - Mesh readers: close() no-ops, MeshDataset.close propagates, include_index_in_metadata, metadata built before pin_memory, _load_extra_boundary_meshes split, multiple-match warning, zarr extra boundary via from_zarr. - record_consumer_stream: to_local() unwrap (CPU and CUDA), non-callable to_local ignored. Domain parallel (multigpu_static): - _plan_key strips sharding shapes; torch's redistribute planner cache stops growing across populated/unpopulated specs. - ShardTensorSpec.__setattr__ drops the cached hash on _sharding_shapes, including the lazy sharding_shapes() population path. - _gather_shard_shapes_for_dim with torch.Size and tensor inputs. - scatter_tensor local shard owns its storage (no view, zero offset). - _convert_args_to_dtensor creates no autograd state below autograd. - _conversion_scope keeps depth attribute at 0; ShardTensor.__new__ is a non-recursive dynamo skip and constructs inside compiled code. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
4250344 to
f69588c
Compare
| return None | ||
| return spawn_generator(self._seed_base, self._epoch, index) | ||
|
|
||
| def _window_indices( |
There was a problem hiding this comment.
This function becomes a generic way to slice into an array, regardless of the type (numpy, torch, zarr, etc) that we can reuse across datareaders.
| metadata = self._get_sample_metadata(index) | ||
| if self.include_index_in_metadata: | ||
| metadata["index"] = index |
There was a problem hiding this comment.
The metadata is moved up here because in PR 4 of the stack, we have an early-exit here for loading domain parallel data, and this lets us not load the meta data in two places.
There was a problem hiding this comment.
(meaning, it's a reorder with no behavior change)
|
|
||
| if self.pin_memory: | ||
| dm = dm.pin_memory() |
There was a problem hiding this comment.
Moving it here just makes MeshReader and DomainMeshReader read meta data and then pin memory in the same order.
| interior=dm.interior, | ||
| boundaries={ | ||
| **dict(dm.boundaries), | ||
| **self._load_extra_boundary_meshes(index), |
There was a problem hiding this comment.
This gets split into two paths deliberately. When we make it domain parallel in PR 4, we have to have a "proto mesh" type object that isn't yet a full mesh, but instead is a mesh-like object that's sharded. This splitring lets us load all the extras, without constructing (and then throwing away) a full mesh later
| def _open_stores( | ||
| self, index: int | ||
| ) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: | ||
| """Open the sample's array stores (async, metadata-only) + attributes.""" |
There was a problem hiding this comment.
this breaks up _load sample into three pieces:
- open stores is attripes, and async opens for arrays (metadata only)
- _load_sample issues the reads itself.
- finalize samples waits on every read issued, or adds missing default values
In the domain parallel component, well make _load_sample_domain_parallel and reuse other components.
| # Per-sample generator: reproducible regardless of read order/thread. | ||
| generator = self._index_generator(index) | ||
| def _open_sample(self, index: int) -> tuple[Any, set[str]]: | ||
| """Open the sample's group and validate required-field availability.""" |
There was a problem hiding this comment.
Same refactor here as tensorstore_zarr to split _load_sample into three pieces (open, read, finalize), but there is an extra move here to push the cyclic index indicis into Reader._window_indices.
| subsample_indices = _cyclic_block_indices( | ||
| array_shape, n_points, generator=generator |
There was a problem hiding this comment.
Here, _cyclic_block_indices is removed, we use the base call now instead.
| def _record(t: torch.Tensor) -> None: | ||
| # Distributed wrapper subclasses (ShardTensor/DTensor) must record | ||
| # their local tensor: record_stream on the wrapper re-enters its | ||
| # dispatch machinery, which has no handling for Stream arguments. | ||
| to_local = getattr(t, "to_local", None) | ||
| if callable(to_local): | ||
| t = to_local() | ||
| if t.is_cuda: | ||
| t.record_stream(stream) | ||
|
|
There was a problem hiding this comment.
This is to ensure we record the stream on the local tensor, if using a sharded tensor, rather than the tensor itself. We don't have record stream machinery in ShardTensor .. that would be a viable path too, I suppose.
| def _plan_key(spec: DTensorSpec) -> DTensorSpec: | ||
| r"""Immutable, canonical key for torch's cached redistribute planner. | ||
|
|
||
| A plain ``DTensorSpec`` hashes and compares on mesh, placements and | ||
| ``tensor_meta`` only -- it carries no lazily-populated sharding shapes, so | ||
| it can sit in a ``functools.cache`` without going stale. | ||
| """ | ||
| meta = spec.tensor_meta | ||
| if meta is not None: | ||
| meta = TensorMeta(tuple(meta.shape), tuple(meta.stride), meta.dtype) | ||
| return DTensorSpec( |
There was a problem hiding this comment.
This changes is all about redistribute caching on specs: the shard shapes was getting added to the spec after it was in the cache table, invalidating the cache, and the cache was growing without limit. We don't need that to route the operations, only to do the operations themselves. SO we actually drop to DTensorSpec for the caching key itself which is totally static.
1 participant
PhysicsNeMo Pull Request
This PR brings standalone fixes and reorganizations that are necessary to start work for the domain-parallel data loading and unified recipe domain parallel support.
Datapipes:
MeshReader/DomainMeshReadergainclose()(fixesMeshDataset.close()raising
AttributeError); the extra-boundary loader is split out.a shared
Reader._window_indices, and a deterministic subsampling windowkey (configured list order instead of set iteration order).
record_consumer_streamunwrapsShardTensor/DTensorleaves.ShardTensor:
scatter_tensorreturns a shard that owns its storage (dynamo could notfakeify the nonzero-offset ranks' views).
mutation (planner cache no longer grows per step).
torch.compilefixes(
_conversion_scope, plain-tensor promotion underno_grad,__new__).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 💬
Coverage note
Codecov patch coverage on this PR is currently an artifact: the PR
Coveragejob runs pytest under--testmon, whose own tracer displaces the outercoverage run, so nothing executed inside tests is recorded (see #1988 for the fix and reproduction). The tests added here will only show up in the patch number once #1988 lands and this branch re-runs CI. Themultigpu_staticadditions run under theci:multi-gpulabel but only count toward coverage in the nightly.