Add a check to ensure shardings are aligned for ops where it matters

[coreyjadams]

PhysicsNeMo Pull Request

This PR targets issue #1943 - Shard tensor can, if users are not careful, accept and happily try to use, sharded tensors where the local sharding is different per tensor yet the global shape is the same. For elementwise operations this would have been sent to DTensor, for example, and silently been wrong in strange patterns, or raised errors in some cases, depending on the local shardings. Either way, not great.

DTensor never checks this, because they have baked in assumptions about how tensors get sharded.

This PR introduces a validation scheme. We keep flexibility, but on ops where the local shapes need to be the same, we can (and do) now apply a validation that checks the local shape is aligned across inputs. I'll annotate the relevant pieces for clarity in the PR.

Description

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

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.

[github-actions]

CODEOWNERS review map

Current for commit 16c0b68ac981. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @coreyjadams — 8 file(s)

• physicsnemo/domain_parallel/init.py
• physicsnemo/domain_parallel/_shard_tensor_spec.py
• physicsnemo/domain_parallel/custom_ops/_tensor_ops.py
• physicsnemo/domain_parallel/shard_tensor.py
• physicsnemo/domain_parallel/shard_utils/attention_patches.py
• physicsnemo/domain_parallel/shard_utils/halo_scatter.py
• physicsnemo/domain_parallel/shard_utils/natten_patches.py
• test/domain_parallel/test_shard_alignment.py

⏳ @negin513 — 8 file(s)

• physicsnemo/domain_parallel/init.py
• physicsnemo/domain_parallel/_shard_tensor_spec.py
• physicsnemo/domain_parallel/custom_ops/_tensor_ops.py
• physicsnemo/domain_parallel/shard_tensor.py
• physicsnemo/domain_parallel/shard_utils/attention_patches.py
• physicsnemo/domain_parallel/shard_utils/halo_scatter.py
• physicsnemo/domain_parallel/shard_utils/natten_patches.py
• test/domain_parallel/test_shard_alignment.py

Comment /codeowners-info to refresh.

[coreyjadams]

Since DTensor has assumptions about sharding shapes, we do an alignment check before passing to DTensor.

[greptile-apps]

Greptile Summary

This PR adds metadata-only validation that rejects mismatched uneven shard boundaries before operations pair local tensor elements.

• Exposes a shared validate_aligned_sharding helper.
• Applies validation to generic DTensor fallbacks and selected cross, attention, NATTEN, and scatter handlers.
• Adds distributed regression tests for aligned and misaligned uneven shards.

Important Files Changed

Filename	Overview
physicsnemo/domain_parallel/_shard_tensor_spec.py	Adds the central alignment validator, but its interpretation of cached cross-sections can produce rank-dependent decisions for repeated shard dimensions.
physicsnemo/domain_parallel/shard_tensor.py	Adds recursive operand-spec collection and invokes alignment validation from both generic fallback paths.
physicsnemo/domain_parallel/custom_ops/_tensor_ops.py	Adds boundary validation before the cross handler selects and executes its local layout.
physicsnemo/domain_parallel/shard_utils/attention_patches.py	Validates K/V and sharded query/mask alignment before positional attention processing.
physicsnemo/domain_parallel/shard_utils/halo_scatter.py	Adds alignment checks for ShardTensor operands in halo-aware scatter handlers.
physicsnemo/domain_parallel/shard_utils/natten_patches.py	Validates Q/K/V boundaries before applying query-derived halo configurations.
test/domain_parallel/test_shard_alignment.py	Covers one-dimensional uneven-shard regressions but not multidimensional repeated-shard layouts.

_{Reviews (1): Last reviewed commit: "Add a check to ensure shardings are alig..." | Re-trigger Greptile}

[coreyjadams]

Reading the local shapes along the sharded dimension directly like this prevents any implicit resolution via collectives, which would be performance hit.

[coreyjadams]

These three lines are the core logic:

• get the current spec's shard shapes
• get the reference spec's shard shapes
• compare

All that's asked here is equality to assure two shard tensors can communicate; sanity checking an individual tensor spec is done at TesnorSpec creation, etc.

[peterdsharpe]

Looks good, thanks for adding the non-regression test!