Geotransolver flare domain parallel

[coreyjadams]

PhysicsNeMo Pull Request

This PR adds support for FLARE and GeoTransolver domain parallelism at the model level. It does not include datapipe support or recipe support yet. Those will come in follow ups.

The models are coming together because FLARE attention is a backend for GeoTransolver; to support GeoT we have to have FLARE anyways so it's cleanest to bring them both together.

There is one tweak to shard_tensor.py incorporated into the PR that needs to be addressed, that really only rears its head in FLARE: when using a wrapper in shard tensor, where we have a torch.Tensor entering the wrapper, it is currently possible for a plain tensor to enter that wrapper. FLARE has this as a parameter that enters the attention mechanism, which is unusual compared to all the other models we see.

To address this, in this PR I add a tweak to fix this in shard tensor: plain torch.Tensor objects will fire promotion (to replicated DTensor, since it's mathematically the same and the machinery is all there) before going to handlers.

New wrappers

There are several new operations in domain_parallel/shard_utils/: a linear_patches.py is the most consequential, and attention_patches.py has substantial additions. A new file grad_ops.py consolidates backward-only operations from several other files (point_cloud_ops.py, conv_patches.py, attention_patches.py) into one logical file and test_grad_ops.py introduces some tests.

Linear Patches

Linear layers are embarrassingly parallel, however pytorch's DTensor op flattens them prior to operating. For some tensors, the view operation itself become problematic especially if the sharded dimension is not the last dimension.
To fix this, the PR adds a custom linear path which routes all linear operations through a local, differentiable op that doesn't hit view issues.

Attention wrappers

The standard attention path in physicsnemo currently is ring attention. FLARE, however, has one tensor which is NOT sharded. This PR adds an attention path for a replicated tensors: it doesn't have to traverse a ring, and so we use a different algorithm.

Layer Updates

For both gale.py and context_projector.py, we have to do a reduction after the tokenization/projection. We duck type, so import shard tensor is unnecessary, and trigger a redistribute. The sharded placement after projecting back out is automatic, just like with transolver.

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.

[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.

[github-actions]

CODEOWNERS review map

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

⏳ @coreyjadams — 13 file(s)

• physicsnemo/domain_parallel/shard_tensor.py
• physicsnemo/domain_parallel/shard_utils/init.py
• physicsnemo/domain_parallel/shard_utils/attention_patches.py
• physicsnemo/domain_parallel/shard_utils/conv_patches.py
• physicsnemo/domain_parallel/shard_utils/grad_ops.py
• physicsnemo/domain_parallel/shard_utils/linear_patches.py
• physicsnemo/domain_parallel/shard_utils/point_cloud_ops.py
• test/domain_parallel/models/harness.py
• test/domain_parallel/models/test_flare.py
• test/domain_parallel/models/test_geotransolver.py
• test/domain_parallel/ops/test_grad_ops.py
• test/domain_parallel/ops/test_linear.py
• test/domain_parallel/ops/test_sdpa.py

⏳ @loliverhennigh — 1 file(s)

• physicsnemo/nn/module/gale.py

No CODEOWNER

• physicsnemo/models/geotransolver/context_projector.py

Comment /codeowners-info to refresh.

[greptile-apps]

Greptile Summary

This PR adds model-level domain parallelism for FLARE and GeoTransolver, including mixed-placement attention, a shard-aware linear path, plain-tensor promotion for registered handlers, and reusable gradient-boundary operations.

• Adds replicated-query/sharded-key-value SDPA and expands mixed-placement attention handling.
• Registers a local shard-aware linear implementation that preserves uneven shard metadata and reduces parameter gradients.
• Resolves reduced GeoTransolver/GALE context tensors before nonlinear attention processing.
• Adds distributed forward/backward coverage for FLARE, GeoTransolver, linear operations, SDPA, and gradient-layout guards.

Important Files Changed

Filename	Overview
physicsnemo/domain_parallel/shard_tensor.py	Extends handler dispatch to promote plain tensor arguments and preserves result/gradient stride metadata during DTensor conversion.
physicsnemo/domain_parallel/shard_utils/attention_patches.py	Adds mixed-placement SDPA, including a distributed log-sum-exp fold for replicated queries over sharded K/V blocks.
physicsnemo/domain_parallel/shard_utils/linear_patches.py	Adds a registered local linear path for non-feature sharding with uneven output-shape bookkeeping and parameter-gradient reductions.
physicsnemo/domain_parallel/shard_utils/grad_ops.py	Consolidates backward-only gradient reduction and contiguous-layout guards used by shard operation patches.
physicsnemo/models/geotransolver/context_projector.py	Resolves tokenizer reductions to Replicate, but the fixed one-element target placement breaks multidimensional DeviceMesh inputs.
physicsnemo/nn/module/gale.py	Resolves slice tokens and context before cross-attention, with the same single-mesh-dimension assumption as the context builder.
test/domain_parallel/ops/test_sdpa.py	Adds forward/backward numerical coverage for both mixed-placement SDPA directions on a one-dimensional mesh.
test/domain_parallel/models/test_geotransolver.py	Adds end-to-end and layer-level domain-parallel GeoTransolver coverage, currently limited to one-dimensional domain meshes.
test/domain_parallel/models/test_flare.py	Adds distributed FLARE coverage for irregular, structured, and spatial-buffer-sharded configurations.

_{Reviews (1): Last reviewed commit: "Update Changelog" | Re-trigger Greptile}

[coreyjadams]

/ok to test cd2316b