ShardTensor: How Per-Atom Fields Flow Across Ranks#

Spatial domain decomposition partitions per-atom tensors across ranks — but without further machinery, every operation on those tensors would be a regular per-rank op with no knowledge of the partition. ShardTensor is the torch.Tensor subclass that carries the partition’s metadata along with the data, and routes select operations through distribution-aware handlers via PyTorch’s __torch_function__ protocol.

This guide is about the mechanism. It complements the distributed user guide (which covers when to use which storage strategy) and Bring Your Own Model: Authoring a Distribution Spec (which covers authoring a spec for a new model).

The subclass approach#

ShardTensor is an “almost transparent” Tensor subclass: it stores the same underlying data buffer as a regular torch.Tensor, plus a small bag of metadata describing the partition:

Field	What it carries	Used by
`_spec`	The `MLIPSpec` governing dispatch behaviour	All op-level handlers
`_meta`	Halo-storage metadata: `n_owned`, `n_padded`, halo routing index	Halo-exchange, halo-correction
`_gather_meta`	Sharded-storage metadata: per-row global IDs, rank assignments	Sharded `index_select`, `scatter_add`
`_config`	The per-rank `ParticleHaloConfig` (process group, ghost width)	All collective ops
`_n_systems`	Number of systems on this rank	Per-system reductions

The wrap is zero-copy: ShardTensor.wrap(t, spec=...) calls t.as_subclass(ShardTensor) and attaches the metadata. The underlying storage is shared; the wrap survives in-place mutations and the autograd graph.

from nvalchemi.distributed.ops import ShardTensor
from nvalchemi.distributed.spec import SPEC_MPNN_HALO

local_positions = torch.zeros(n_padded, 3, device="cuda")
shard = ShardTensor.wrap(
    local_positions,
    spec=SPEC_MPNN_HALO,
    meta=halo_meta,         # ParticleHaloMetadata describing this rank's slice
    config=halo_config,     # ParticleHaloConfig with the process group
    n_systems=n_systems,
)

Most code never constructs ShardTensor directly: DistributedModel promotes data.positions (and data.charges if present) to ShardTensor in its halo-storage call path before invoking the wrapper, so per-atom ops inside the wrapper’s forward see a ShardTensor and dispatch accordingly.

`__torch_function__` dispatch#

Every torch op called on a ShardTensor runs through ShardTensor.__torch_function__(func, types, args, kwargs) — the standard subclass-hook PyTorch provides. The dispatch is predicate-based: a small registry of handlers, each tagged with a predicate (func, args, kwargs) -> bool, is consulted in order. The first matching handler runs; if none match, the op falls back to the default torch.Tensor.__torch_function__.

Dispatch decision tree for an op called on a ShardTensor.#

Three handler families are registered today:

Halo correction. Fires on scatter_add_ / index_add_ calls where the destination is a ShardTensor carrying halo metadata. After the local scatter, the handler does halo_reverse_exchange + halo_forward_exchange so halo rows contribute their partial sums back to owners and the halo is re-populated with the corrected owner values for downstream ops.
Per-system reduce. Fires on scatter_add_ calls whose target is a per-system buffer (shape (n_systems, F)) and whose source is per-atom. Slices halo rows off the source first (so each atom contributes once), then scatters locally and all-reduces across the mesh.
Distributed scatter / index_select. Fires for sharded-storage models. Routes the op via global IDs: an index_select with cross-rank target rows gathers them via all_to_all_v; a scatter_add with cross-rank source rows likewise.

The registry is in nvalchemi.distributed._core.shard_tensor and is keyed by op + predicate so multiple handlers can coexist for the same op (e.g. halo-correction for one shape, per-system-reduce for another).

Halo storage in detail#

Halo-storage layout across two ranks for a 6-atom system.#

At step start, each rank’s halo rows are stale. The halo exchange populates them by all-to-all-v of owned-row data into the partner ranks’ halo slots, with _meta.halo_routing carrying the index table. For the duration of the step, every read of positions[j] where j is a halo row resolves to a current copy of rank r(j)’s owned atom — so cross-rank pair distances are computed locally with no further communication.

Halo writes are different. When a model writes into a halo row via out.scatter_add_(0, receiver, msg) and receiver[e] happens to be a halo atom, the write only contributes a partial sum on this rank. The corresponding owner on the other rank holds its own partial sum from its own edges. Halo correction reverses this: after the scatter, halo-row partial sums are routed back via halo_reverse_exchange and added to the owner’s value; then halo_forward_exchange repopulates this rank’s halo with the combined owner result so downstream ops in the same forward pass see consistent per-atom features.

The handler is registered on scatter_add_ / index_add_. A model author who writes a standard PyTorch MPNN

out = torch.zeros_like(x)
out.scatter_add_(0, receivers, msg)

gets halo correction for free iff out is a ShardTensor — which it is automatically when x is, because torch.zeros_like(x) propagates the subclass.

Sharded storage in detail#

Sharded-storage models hold only owned rows on each rank — there’s no halo. Cross-rank lookups happen on demand. ShardedBatch’s _atom_fields carry per-row global-ID metadata; when a model does x.index_select(0, idx) with idx containing cross-rank global IDs, the dispatch handler:

Inspects idx against the _gather_meta.rank_assignment table to figure out which rank owns each requested row.
Issues an all_to_all_v to ship the requested rows to this rank.
Reorders the result back into the order idx requested.

The reverse holds for scatter_add with cross-rank receivers: locally-grouped contributions are shipped to the owner ranks where the actual scatter occurs.

_gather_meta carries a n_global sentinel for “this slot is padding” — out-of-range indices in idx resolve to a known empty contribution rather than triggering a CUDA out-of-bounds assertion.

Subclass propagation guarantees#

Two PyTorch behaviours that the framework relies on:

Like-shaped allocator ops preserve subclass. torch.zeros_like(x), torch.empty_like(x), x.new_zeros(...), etc. return a ShardTensor when x is a ShardTensor, with the same _spec / _meta / _config. This is what makes the toy MPNN pattern in Bring Your Own Model: Authoring a Distribution Spec work without explicit wrap calls inside the model body.
Most ops downcast to torch.Tensor. x[i], x + y, linear(x) — these go through the default __torch_function__ path which produces a plain Tensor view of the underlying storage. The autograd graph flows through this view; the ShardTensor subclass identity is dropped. This is fine for ops that don’t need cross-rank communication.

The split between “preserves subclass” and “drops subclass” is deliberate. Halo-correction needs the destination of scatter_add_ to be a ShardTensor (so it can reach the metadata); intermediate features after a Linear layer don’t, because Linear is a per-rank op.

Custom ops and `OpAdapter`#

Warp / Triton / Numba / generic CUDA kernels wrapped as @torch.library.custom_op are opaque to subclass dispatch: the kernel does wp.from_torch(t) (or the equivalent) internally, bypassing __torch_function__. Without further help, calling such an op on a ShardTensor would unwrap to a plain Tensor (losing the metadata) and run the kernel as if the input were single-process.

OpAdapter declares the distribution semantics for one such kernel:

from nvalchemi.distributed.ops import GatherInputs, ScatterOutputs
from nvalchemi.distributed.spec import OpAdapter

OpAdapter(
    op=torch.ops.mymodel.fused_kernel.default,
    arg_transforms={0: GatherInputs()},     # halo-pad input position 0
    output_transforms={0: ScatterOutputs()}, # halo-correct output position 0
)

The adapter goes on the spec’s distribution.custom_ops. At scope-entry the framework’s AdapterRegistry walks the spec, installs a ShardTensor handler on each op handle, and the kernel becomes distribution-aware: when called with a ShardTensor input, the handler runs the declared arg_transforms, calls the kernel on plain tensors, then runs the declared output_transforms and re-promotes outputs.

The available transforms are:

Transform	Pre-/post-kernel action
`GatherInputs`	halo-pad an owned input to `(n_padded, *F)`
`GatherInputsFull`	full-gather a sharded input to `(n_global + 1, *F)`
`SliceOwned`	slice a halo-padded input to `(n_owned, *F)`
`ScatterOutputs`	`halo_reverse + halo_forward` on a per-atom output
`AllReduceSum`	cross-mesh `SUM` all-reduce on a partial output
`SliceOutputsOwned`	slice an `(n_global + 1, F)` output back to `(n_owned + 1, F)`

See Bring Your Own Model: Authoring a Distribution Spec for an end-to-end OpAdapter authoring example with a Warp kernel.

When you don’t need ShardTensor#

If your wrapper’s forward is built entirely from torch ops that the framework already handles (scatter_add_, index_select, scatter_add, etc.) and you stay within a single storage strategy, you generally don’t touch ShardTensor directly. The framework promotes data.positions and the subclass propagation handles the rest.

You do need ShardTensor when:

Your model has a per-layer node-feature buffer (e.g. message-passing state) that scatter writes target. Wrapping that buffer once via ShardTensor.wrap(...) in adapt_input is enough — see the MACE wrapper’s adapt_input for the canonical pattern.
You’re authoring a custom op via OpAdapter and need to declare what shape the kernel expects and produces.

You don’t need ShardTensor for:

Pure per-atom ops with no aggregation (per-atom MLP, embeddings).
Ops on per-system tensors (scatter_add_ with target shape (n_systems, F) is automatically routed via per-system-reduce when system_reductions=True on the spec).

Reference: where ShardTensor lives#

The full implementation is in nvalchemi.distributed._core.shard_tensor. The upstream-candidate boundary linter (tools.check_core_imports) keeps this module chemistry-vocabulary-free; it’s the basis for any future upstream contribution to PhysicsNeMo or related projects.

Next steps#

Bring Your Own Model: Authoring a Distribution Spec walks through declaring a spec for a new wrapper, authoring an OpAdapter for a Warp kernel, and using trace_and_validate to confirm distributed correctness.
The runnable examples in examples/distributed/04_* and examples/distributed/05_* exercise both patterns end-to-end.