Avoid shape-negotiation all_to_all in ShardTensor redistribute when shapes are known (ShardTensor Optimization #1) - #1779
Conversation
Redistribute's Shard->Shard transpose negotiates recv-buffer shapes with a collective every call, even though current_spec already records each rank's shard shape whenever it was built with known (non-"infer") sharding shapes. Passing current_spec through is prep for computing those shapes analytically instead of over the network.
_to_new_shard_dim ran a full all_to_all just to learn recv-buffer shapes before the real data-movement all_to_all -- doubling the collective count for every Shard->Shard transpose. When current_spec already has known per-rank shard shapes (and they're still fresh, i.e. not stale from an earlier hop in the same multi-hop redistribute call), derive recv shapes analytically instead: the target-dim chunk size is rank-independent (every rank computes it the same way), and the current-dim extent per sender is already on current_spec. Falls back to the original negotiation otherwise. Fixes a 2D-mesh multi-hop staleness case (test_shard_tensor_redistribute2d S3+R) by verifying local_tensor's actual shape still matches what's recorded before trusting it.
|
Nice! |
Greptile SummaryThis PR eliminates a shape-negotiation
Important Files Changed
Reviews (1): Last reviewed commit: "Apply ruff format to the redistribute al..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
coreyjadams
left a comment
coreyjadams
left a comment
There was a problem hiding this comment.
This is, I think, the most complicated of the redistribute paths in ShardTensor. with dynamic shapes along multiple axes, this can get really complicated. I 1000% see the performance benefit of the fast path here and you are right to push for it, and I agree. What do you think about getting a little more aggressive on the testing for this particular op, and covering more scenarios? I think our static tests are completing in only a few minutes currently, we have a ton of runtime budget for more tests in multigpu-static and this is a great place to add logic checking.
… shapes Address review: the recv-shape fast path was gated on a rank-local comparison (recorded spec shape vs local_tensor.shape), which can diverge across ranks (e.g. empty shards under uneven chunking) and mismatch collectives. The staleness signal is now an explicit flag threaded from the redistribute hop loop -- true only until the first transform mutates the local tensors -- which depends only on the hop sequence and is therefore identical on every rank. The shape comparison is kept as a loud error for corrupt specs. Also fixes compute_sharding_shapes_from_chunking_global_shape to record each rank's true chunk size along the varying mesh dim; it previously wrote this rank's chunk size into every per-rank entry, which was only correct for evenly divisible sharding. Adds multigpu-static coverage: even and uneven sharding for the 1D and multi-hop 2D transpose paths, empty-shard redistribution, and a fast-path vs negotiated-fallback equivalence test.
negin513
changed the title
|
Re-ran the isolated redistribute benchmark against the final branch state (
|
Per review: use the available runtime budget for more aggressive logic checking on the most complicated redistribute path. - Fast-path vs negotiated-fallback bitwise equivalence on 2D multi-hop paths (pure transpose, transpose chain, transpose-then-gather), with both tensor dims rank-dependent, even and uneven. - Empty shard along one mesh dim combined with uneven sharding along the other, through the transpose chain and double-gather paths. - Backward through the transpose: fast path and fallback must produce bitwise-identical gradients (covers the autograd reverse path). - Adds a skip-marked regression test documenting a pre-existing crash (reproduced on main @ 8e76840): resharding onto a tensor dim with extent smaller than the mesh size segfaults in the all_to_all.
|
Resolves NVIDIA/physicsnemo-roadmap#2761 and NVIDIA/physicsnemo-roadmap#828 | relevant for HealDA (NVIDIA/physicsnemo-roadmap#2172) |
|
Added more tests in 356c400. @coreyjadams I think this should be ready for re-review. |
|
/blossom-ci |
|
/ok to test 356c400 |
|
Disabling auto merge exclusively so we can refresh the CI label and run multi-gpu CI on this. Looks good to go, let's make sure it passes the CI here before it merges. :) |
|
Once another PR lands, and we update this one, it will refresh and we can queue it. |
|
/ok to test f522d04 |
2 participants
Summary
TLDR: avoid an extra
all_to_allwhen shapes are know! 2.5x speed-up.While working on merging HealDA, we realized
redistribute'sShard->Shardtranspose does twoall_to_allcollectives instead of one -- the first just asks other ranks "what shape is your chunk?" before receiving anything.... The overhead from the extraall_to_alladds up fast for cases where we reshard repeatedly, such as HealDA where we do two reshards per block.In most cases it doesn't need to ask at all: each rank's own chunk size is deterministic, and the sender's shape is already sitting on
current_specfrom when the tensor was constructed.This PR calculates the recv-buffer shape locally instead of asking over the network: the target-dim chunk size is deterministic (every rank computes its own the same way), and the sender's current-dim extent is already recorded on
current_specfrom when the tensor was constructed. It only falls back to the original negotiation when that info genuinely isn't available.Why the staleness check is added?
current_specis fixed for the whole (possibly multi-hop)redistributecall, but a 2D mesh can chain multiple transform steps in one call (e.g.Shard/Shard -> Shard/Replicategathers one mesh dim, then transposes the other). An earlier hop can already have reshapedlocal_tensorby the time a later hop runs, making the recorded shapes describe a state that no longer exists. Added a check thatlocal_tensor.shapestill matches whatcurrent_spechas on record before trusting the shortcut -- caught bytest_shard_tensor_redistribute2d'sS3+Rcase, which failed with a shape mismatch before this guard was added.Compatibility
No caller changes needed -- this helps existing code as-is, including
sharding_shapes="infer"(its default mode already records the needed shapes eagerly).Performance
Isolated the redistribute primitive (4xGB200, T=16, X=12288, C=1536, bf16, even sharding, caller already passing known
sharding_shapesso this measures only this fix's effect):Test plan
torchrun --nproc_per_node=4 -m pytest test/domain_parallel/test_redistribute.py --multigpu-static -q-- 11 passed (1D/2D mesh, even/uneven sharding), 11 skipped (dynamic-mode variants)torchrun --nproc_per_node=4 -m pytest test/domain_parallel/ --multigpu-static -q-- confirmed identical set of 8 pre-existing, unrelated failures with and without this changepre-commit(ruff check, ruff format, interrogate, license, large-file check) -- all pass on the changed fileResolves NVIDIA/physicsnemo-roadmap#2761
Related
torch.compilesupport) -- that PR doesn't touch_shard_redistribute.py/_shard_tensor_spec.py, so this applies cleanly on top either wayglobal_shapeonfrom_localto unlock its existing no-comm"chunk"path) -- that addresses a separateall_gatherinfrom_local; this PR's fast path only activates once callers can avoid"infer", which that follow-up enables through the public APIChecklist
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.