Expose global_shape on ShardTensor.from_local to unlock the no-comm chunk path [ShardTensor Optimization 2]

[negin513]

PhysicsNeMo Pull Request

Description

from_local defaults to sharding_shapes="infer", which pays a real all_gather on every call to learn each rank's shard shape -- even when the shard is evenly divided and the shape is derivable locally. The internal spec builder (_infer_shard_tensor_spec_from_local_chunks) already supports a communication-free "chunk" mode, but it requires a global_shape argument that from_local never exposed -- so that mode was unreachable from the public API (passing sharding_shapes="chunk" always raised ValueError: If sharding_shapes is 'chunk', global_shape must be provided, with no way to satisfy it).

This PR threads an optional global_shape parameter through from_local -> _FromTorchTensor.forward -> the existing internal helper, and grows _FromTorchTensor.backward's return tuple by one None for the new non-differentiable input. The codebase already uses this exact pattern internally (_resolve_spec_for_dtensor calls the helper with sharding_shapes="chunk", global_shape=dtensor.shape) -- this just makes it available to from_local callers.

Fully additive: global_shape defaults to None everywhere, so existing "infer" and dict callers are unaffected. _FromTorchTensor.apply has exactly one caller (from_local itself), so nothing else needed updating.

Closes #1780
Resolves NVIDIA/physicsnemo-roadmap#2719

Testing

• New-behavior checks on 4xGB200 (torchrun --nproc_per_node=4): "chunk"+global_shape produces a spec identical to "infer"'s (same global shape, same per-rank sharding shapes), round-trips through redistribute, flows gradients through the 5-arg autograd function, still raises ValueError for "chunk" without global_shape, and the default no-kwarg path is unchanged.
• torchrun --nproc_per_node=4 -m pytest test/domain_parallel/ --multigpu-static -q -- 394 passed, 0 failed on top of the merged ShardTensor Refactor #1556 refactor.
• pre-commit (ruff check/format, interrogate, license) clean on the changed file.

Related

• Companion to Avoid shape-negotiation all_to_all in ShardTensor redistribute when shapes are known (ShardTensor Optimization #1) #1779: that PR removes the doubled all_to_all in redistribute's Shard->Shard transpose; this one removes the all_gather in from_local. Both are independent redundant collectives found while profiling the context-parallel reshard benchmark in HealDA v2 Architecture #1758, and they stack.
• Related to NVIDIA/physicsnemo-roadmap#2712 and NVIDIA/physicsnemo-roadmap#2607.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

None.

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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR threads an optional global_shape parameter through ShardTensor.from_local → _FromTorchTensor.forward → the existing _infer_shard_tensor_spec_from_local_chunks helper, making the pre-existing communication-free "chunk" spec path reachable from the public API for the first time. Previously, sharding_shapes="chunk" always raised a ValueError because from_local never exposed global_shape.

• global_shape defaults to None everywhere, so all existing "infer" and dict callers are unaffected.
• _FromTorchTensor.backward correctly extends its return tuple from 4 to 5 values to match the new 5-argument forward signature, maintaining autograd correctness.
• The _FromTorchTensor class-level docstring is slightly stale (still claims comm is always required); a minor inline fix is suggested.

Important Files Changed

Filename	Overview
physicsnemo/domain_parallel/shard_tensor.py	Threads an optional global_shape parameter through from_local → _FromTorchTensor.forward → _infer_shard_tensor_spec_from_local_chunks, enabling the communication-free "chunk" spec path. backward return tuple is correctly extended to 5 values to match the 5 forward inputs. Change is fully additive; only the class-level docstring is slightly stale.

Comments Outside Diff (1)

1. physicsnemo/domain_parallel/shard_tensor.py, line 514-515 (link)

   The _FromTorchTensor class docstring still claims that global shape is always inferred via collective communication, but that statement is now incorrect. With global_shape provided and sharding_shapes="chunk", the forward pass is entirely communication-free. A reader of this class-level docstring will have a misleading mental model before they ever reach the method signatures.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "Expose global_shape on ShardTensor.from_..." | Re-trigger Greptile}

[coreyjadams]

Thanks, looks reasonable!

[negin513]

Perf numbers for the new path (4xGB200, torchrun --nproc_per_node=4, local shard (1,4,12288,1536) bf16, Shard(1) on a 1D mesh, 50 iters after warmup):

from_local mode	ms/call
sharding_shapes="infer" (default, pays all_gather)	0.378
sharding_shapes="chunk" + global_shape (this PR)	0.036

~10x faster wrapper construction once the collective is gone. In the HealDA reshard pattern (#1758) from_local is called twice per block, so this stacks with #1779's redistribute savings.

[negin513]

Here are the numbers:

Reshard round trip	fwd	fwd+bwd
main (today)	2.661 ms	5.532 ms
#1779 only	1.471 ms	3.233 ms
#1791 only	1.806 ms	4.943 ms
combined	0.768 ms	2.602 ms

[negin513]

/ok to test 8d74adc