Distributed (spatial domain decomposition)#

The nvalchemi.distributed package runs the toolkit’s dynamics and models across multiple GPUs by partitioning atoms in space. This page is organised around the tasks a developer actually performs — running an existing model distributed, bringing your own model under domain decomposition, and writing the wrapper code that makes a model distribution-aware — rather than as a flat symbol list. For the concepts behind each piece, read the companion guides first: Distributed Simulations (overview and the two storage strategies), Bring Your Own Model: Authoring a Distribution Spec (writing a wrapper + spec), and ShardTensor: How Per-Atom Fields Flow Across Ranks (the sharded-tensor dispatch layer).

Everything below is importable from the package root, e.g. from nvalchemi.distributed import DomainParallel; the few symbols that live in a submodule (the SPEC_* presets, StrategyKind) are noted where they appear.

Running a model under domain decomposition#

The entry point is DomainParallel: wrap any single-GPU BaseDynamics integrator or optimizer with a DomainConfig, and it partitions the system, exchanges halos, runs the model, consolidates outputs, and migrates atoms across domain boundaries each step. The model wrapper, hooks, and integrator are unchanged from the single-process API — the only additions at the user layer are the config and the wrap.

DomainConfig carries the three concerns a scope needs: the interaction cutoff/skin and ghost_width (halo geometry), the process-mesh topology (mesh/mesh_dim/grid_dims), and the strategy — HALO (spatial domain decomposition, the default) or GRAPH_PARTITION (a node partition for models that build their own neighbour list). Set compile=True to let the framework own a shape-stable compiled forward. Hooks run at a HookScope (LOCAL per-rank, GLOBAL after an all-gather, or RANK_ZERO).

Under the hood, DistributedModel (or DistributedPipelineModel for a composed pipeline) is the per-step adapter that owns halo exchange, neighbor rebuild, and output consolidation; ShardedBatch is the partitioned per-atom state it operates on, produced by SpatialPartitioner. Most users never touch these directly, but they are the seams a custom runtime hooks into.

`DomainParallel`	Wraps any `BaseDynamics` subclass with spatial domain decomposition.
`DomainConfig`	Configuration for one spatial domain-decomposition scope.
`HookScope`	Determines which ranks execute a hook callback.
`DistributedModel`	Wrap an atomic single-process model wrapper for domain-decomposed inference.
`DistributedPipelineModel`	Domain decomposition of a composed (pipeline) model.
`ShardedBatch`	A `Batch` distributed across a 1-D `DeviceMesh`.
`SpatialPartitioner`	Assigns atoms to spatial sub-domains on a Cartesian grid.

StrategyKind

Which parallelization strategy a distributed scope runs under.

Bringing your own model: the distribution spec#

A model tells the framework how to parallelize it by returning a MLIPSpec from BaseModelMixin.distribution_spec(strategy). The spec is the single source of truth for how each output is combined, which storage policy applies, and which opaque kernels need adapters. Declare each output once with an OutputSpec (see Declaring model outputs); the spec lowers that onto the wire fields and round-trips through to_dict() / from_dict(). Specs compose with | so a model can start from a preset and override a field.

Most models never write a spec by hand — the SPEC_* presets in nvalchemi.distributed.spec cover the shipped families (SPEC_MPNN_HALO for scatter-heavy message-passing nets, SPEC_LJ_HALO, SPEC_UMA_HALO, SPEC_EWALD_HALO, SPEC_PME_HALO, SPEC_DFTD3_HALO, and SPEC_MPNN_GP for the graph-partition strategy). Start from the preset that matches your model’s communication pattern and adjust. DistributionSpec is the framework-generic layer underneath (storage policy + custom-op/third-party-helper tuples); CompilePolicy and ForceStrategy tune the compiled-forward and force-derivation behaviour.

`MLIPSpec`	What an MLIP needs from the distributed framework.
`DistributionSpec`	Framework-generic distributed spec.
`CompilePolicy`	How a model wants `torch.compile` driven under domain decomposition.
`ForceStrategy`	How a model's forces are produced under a distributed forward.

`SPEC_MPNN_HALO`	MACE, NequIP, Allegro, ORB.
`SPEC_LJ_HALO`	Lennard-Jones pair potential.
`SPEC_UMA_HALO`	UMA (eSCN-family) via the halo storage policy, with fairchem graph parallel disabled.
`SPEC_EWALD_HALO`	halo storage.
`SPEC_PME_HALO`	halo storage.
`SPEC_DFTD3_HALO`	halo storage, no global coupling.
`SPEC_MPNN_GP`	atoms split by a balanced index range (no spatial halo), each rank owning the edges into its nodes.

Adapters for opaque kernels#

Domain decomposition works by dispatching per-atom tensor operations through a ShardTensor (see ShardTensor: How Per-Atom Fields Flow Across Ranks). Kernels that bypass __torch_function__ — Warp/Triton launches, @torch.jit.script ops, or a model’s internal graph builder — are invisible to that dispatch and must be declared as adapters on the spec’s custom_ops. Pick the adapter that matches how the kernel is invoked: OpAdapter (a registered custom op), MethodAdapter (a method on a named class), FunctionAdapter (a module-level function), PythonAdapter (an arbitrary attribute patch), or JitAdapter (a scripted op needing marshalling across the ShardTensor boundary). AdapterRegistry collects them and AdapterStatus reports whether each applied.

`OpAdapter`	Adapt a `@torch.library.custom_op` / `@torch.library.triton_op` kernel to ShardTensor-aware dispatch.
`MethodAdapter`	Wrap a class method: intercept the call, transform an argument, then invoke the original — as opposed to `PythonAdapter` / `JitAdapter`, which replace a module-level function outright.
`FunctionAdapter`	Adapt a module-level function named by the real function object.
`PythonAdapter`	Replace a plain-Python module-level helper with a distributed-aware version.
`JitAdapter`	Replace a `@torch.jit.script`-decorated module-level helper so a ShardTensor can cross it safely on the distributed path.
`AdapterRegistry`	Owns the install / restore lifecycle for a set of adapters.
`AdapterStatus`	Introspectable record of one adapter's lifecycle state.

Declaring model outputs#

Consolidation needs to know, for every model output, its shape (per-atom vs per-system) and how each rank’s partial value is combined into the global result. Declare this with an OutputSpec per output on the model’s MLIPSpec:

outputs={"stress": OutputSpec(kind=OutputKind.PER_GRAPH, reduce=Reduce.ALL_REDUCE)}

OutputKind covers the shape axis (PER_NODE / PER_GRAPH, plus GLOBAL passthrough and UNKNOWN fallback) and Reduce the combine rule (NONE per-kind default, ALL_REDUCE to sum partials across the mesh, or OWNED_ONLY for values already correct on every rank). Getting these right is what makes a distributed forward numerically match the single-process result.

`OutputSpec`	How one named model output is shaped and combined under DD.
`OutputKind`	Per-output classification used by consolidation.
`Reduce`	How an output's per-rank value is combined into the global value.

Validating a spec#

Before trusting a new wrapper, run trace_and_validate(): it traces the model under domain decomposition, checks that every per-atom operation and opaque kernel is covered, and returns a report with a pass/fail verdict, the applied fixes, and per-layer diagnostics pinpointing any uncovered op. This is the first thing to reach for when a distributed forward disagrees with the single-process baseline.

trace_and_validate

Infer a distribution spec, validate it on a single-GPU multi-process run, and (optionally) auto-fix when validation fails.

Writing adapter bodies: the intent vocabulary#

Inside a distributed method or adapter body, express what you need by intent rather than naming the mechanism, and the helper does the right thing under the halo policy, in single-process, and under torch.compile. Mark a wrapper method as distribution-aware with distributed_method(); then call refresh_neighbors() to update halo rows after a neighbor rebuild, scatter_to_owners() to reverse-sum halo contributions back to owning ranks, system_sum() to reduce a per-system quantity across the mesh, and to_local() / localize() to drop to owned-only rows. current_dd_context() exposes the live context, autograd_target() returns the in-graph leaf to differentiate against, and Scope selects owned vs padded extent.

`distributed_method`	Decorate a `MethodAdapter` body that only diverges under domain decomposition.
`refresh_neighbors`	Populate this rank's neighbor (ghost) rows of a per-node tensor.
`neighbor_refresh_adapters`	Build adapters that recombine each module's per-node `forward` output across ranks.
`scatter_to_owners`	Fold per-edge contributions written into ghost rows back to owners.
`system_sum`	Sum per-node values into per-system totals, without double-counting.
`to_local`	Return the plain local tensor backing a ShardTensor, else `x` unchanged.
`localize`	Run `to_local()` over every value in a model-input dict.
`current_dd_context`	Return the live DD context for the current forward.
`autograd_target`	Return the tensor to pass as `torch.autograd.grad()`'s `inputs=`.
`Scope`	Which rows a per-system reduction sums, and whether it crosses ranks.

Graph-parallel padding#

Under StrategyKind.GRAPH_PARTITION (and any compiled forward), per-rank atom and edge counts drift as the system evolves, which would force recompilation every step. The GraphPadder family pads counts to stable caps so the compiled graph is reused: COOPadder for sparse edge indices, DensePadder / DenseBatchPadder for dense neighbor matrices, with resolve_cap() choosing the padded size.

`GraphPadder`	How a model's graph representation is padded to a fixed capacity.
`COOPadder`	Built-in `GraphPadder` for COO `edge_index` graphs.
`DensePadder`	Built-in `GraphPadder` for dense `(N, K)` neighbor-matrix graphs.
`DenseBatchPadder`	Built-in `GraphPadder` for dense `(N, K)` neighbor-matrix `Batch`es (AIMNet2).

resolve_cap

Grow-only fixed-shape capacity for key, >= real + extra.

Halo and resharding primitives#

Lower-level building blocks for custom distributed flows: ParticleHaloConfig configures the halo exchange, and reshard_by_destination() moves atoms to new owning ranks (the atom-migration primitive DomainParallel uses each step).

ParticleHaloConfig

Configuration for particle-based halo exchange.

reshard_by_destination

Redistribute tensor elements to new ranks based on per-element destinations.