nvalchemi-distributed

nvalchemi Distributed (Domain Decomposition)

Overview

Domain decomposition (DD) splits one atomic system across several GPUs so a simulation that doesn't fit — or doesn't run fast enough — on a single card can scale out. The same model wrapper, hooks, and integrators you use single-process run unchanged; you add one wrapper around each.

• {class}~nvalchemi.distributed.DomainParallel wraps any {class}~nvalchemi.dynamics.base.BaseDynamics integrator/optimizer and drives it across the mesh.
• Under the hood it wraps the model in a DistributedModel, which reads the model's distribution_spec to know how to shard/gather each field.

from nvalchemi.distributed import (
    DistributedManager, DomainConfig, DomainParallel, HookScope,
)

Launch with torchrun (or SLURM): DD is one process per GPU.

torchrun --nproc_per_node=4 my_distributed_md.py

There are four things you may need to do. Pick the section you need:

1. Run a shipped model under DD → §1
2. Choose halo vs graph-partition → §2
3. Make your own model run under DD (author a distribution_spec) → §3
4. Write a custom integrator that stays correct under DD → §4

1. Run a domain-decomposed model

Bootstrap the process group + mesh with DistributedManager, wrap the model as usual, then wrap the integrator in DomainParallel. Build the full system on rank 0, partition() it, and run().

Key points:

• partition() takes the full system on rank 0 (None elsewhere) and returns each rank's owned Batch. run() loops step() for n_steps.

• with DomainParallel(...) makes teardown exception-safe; keep the process-group lifecycle (initialize / cleanup) at launcher scope.

• The initial system (build_full_system, built on rank 0) is a Batch whose per-atom fields must include everything the integrator reads. For MD that means positions, atomic_numbers, atomic_masses, cell, pbc, and velocities. Masses/positions/cell/pbc are AtomicData(...) constructor args, but velocities are attached separately — data.add_node_property("velocities", v) — then Batch.from_data_list([data]). NVTLangevin / NVE read batch.velocities + batch.atomic_masses and crash at step 0 if they're absent:

  data = AtomicData(positions=pos, atomic_numbers=z, atomic_masses=m,
                    cell=box.unsqueeze(0), pbc=pbc)     # (masses in the ctor)
  data.add_node_property("velocities", v)               # velocities added after
  full = Batch.from_data_list([data], device=device)

• Hook placement: neighbor-list / compute hooks go on the inner integrator (they fire on the padded compute view); trajectory/logging hooks go on the outer DomainParallel with HookScope.RANK_ZERO (gather to rank 0) or HookScope.GLOBAL (gather to every rank).

Runnable references: examples/distributed/03_mace_nvt_distributed.py (NVT), 06_mace_npt_distributed.py (NPT), 07_fire_nvt_dd.py (FIRE + 2-D pipeline×DD).

2. Choose the strategy: halo vs graph-partition

Two strategies ship, selected with DomainConfig.strategy ({class}~nvalchemi.distributed.config.StrategyKind). The default is halo.

from nvalchemi.distributed.config import StrategyKind
 
# Halo (default):
cfg = DomainConfig(cutoff=wrapper.cutoff, skin=0.5, mesh=mesh)
# Graph-partition:
cfg = DomainConfig(cutoff=wrapper.cutoff, skin=0.5, mesh=mesh,
                   strategy=StrategyKind.GRAPH_PARTITION)

Rule of thumb: use halo for a short-range MPNN with a real cutoff (it scales to large N as a capacity play); use graph-partition for a model that rebuilds its own graph and can't take a pre-padded view, or to fit a single system past one GPU's memory.

Shipped presets (from nvalchemi.distributed.spec) — a wrapper declares one of these as its distribution_spec (§3):

Note: SPEC_LJ_HALO is for the shipped LJ's opaque kernel, not a BYO pair potential — if you write a differentiable pair potential in plain PyTorch, its forces come from autograd and the correct preset is SPEC_MPNN_HALO.

Compiled runs: add compile=True to the DomainConfig. The forward is fixed-shape (padded to per-rank caps), so after a short warm-up the trajectory runs recompile-free — MD-ready. Supported by the MACE (incl. cuEquivariance), AIMNet2, and UMA wrappers.

3. Add a distribution_spec (bring your own model)

To make an arbitrary wrapped model (see nvalchemi-model-wrapping) run under DD, give it a distribution_spec that returns an {class}~nvalchemi.distributed.spec.MLIPSpec. The framework reads it — the wrapper's forward stays free of any distributed code. Four steps:

(1) Wrap your model with BaseModelMixin as usual.

(2) Declare a spec. If your model is a scatter-heavy MPNN with autograd forces, return a preset:

from nvalchemi.distributed.spec import SPEC_MPNN_HALO
 
class MyWrapper(nn.Module, BaseModelMixin):
    def distribution_spec(self, strategy=None):
        return SPEC_MPNN_HALO

Getting the wrapper right under halo. For an autograd/scatter model that returns SPEC_MPNN_HALO, four things in your forward / adapt_input decide correctness — miss them and you get silently wrong energies, not a crash:

• Emit per-atom energies over ALL local atoms (owned + ghost). Return an (n_local,) atomic-energy tensor and let the framework reduce ghost rows; don't pre-slice to owned.
• Rebind in-place scatters: write ae = ae.scatter_add_(0, recv, e_edge) and use the returned tensor — under DD the dispatch returns a new, cross-rank halo-corrected tensor; the pre-rebind value is wrong.
• Filter neighbor-list sentinels in adapt_input: halo padding can leave receiver indices ≥ n_atoms; drop them with valid = (edge_index[0] < n) & (edge_index[1] < n).
• Periodic systems: compute displacements with the minimum-image convention from the cell (frac = rij @ inv_cell; frac -= frac.round(); rij = frac @ cell). It is a no-op on halo ghost edges (the ghost already sits at its image) and correct single-process, so no explicit shift vectors are needed.

Your wrapper's ModelConfig should declare forces as autograd — autograd_outputs=frozenset({"forces"}), autograd_inputs=frozenset({"positions"}), neighbor_config=NeighborConfig(cutoff=…, format=NeighborListFormat.COO). See nvalchemi-model-wrapping and examples/distributed/04_byo_pytorch_mpnn.py for the full wrapper.

If the model calls a custom op (a Warp/Triton kernel, a fused scatter) the tracer can't see through, build the spec explicitly and declare per-op behaviour with an adapter:

import torch
from nvalchemi.distributed.ops import HaloStoragePolicy, ScatterOutputs
from nvalchemi.distributed.spec import (
    DistributionSpec, MLIPSpec, OpAdapter, OutputKind,
)
 
def _build_spec():
    return MLIPSpec(
        distribution=DistributionSpec(
            policy=HaloStoragePolicy(),        # halo layout
            custom_ops=(                       # one OpAdapter per opaque kernel
                OpAdapter(
                    op=torch.ops.mypkg.my_energy_op.default,
                    arg_transforms={},                       # inputs already halo-padded
                    output_transforms={0: ScatterOutputs()}, # halo-correct output 0
                ),
            ),
        ),
        output_kinds={"energy": OutputKind.PER_GRAPH,
                      "forces": OutputKind.PER_NODE},
    )

• policy — storage layout: HaloStoragePolicy (halo) or GraphParallelPolicy (node-partition). output_kinds tags each named output's shape (PER_NODE, PER_GRAPH, GLOBAL) so consolidation knows whether to halo-reverse, all-reduce, or pass it through.
• Adapters declare how a black-box op distributes. OpAdapter(op=, arg_transforms={i: …}, output_transforms={i: …}) goes in DistributionSpec.custom_ops; MethodAdapter (swap a third-party module method) and PythonAdapter / JitAdapter (replace a helper) go in third_party_helpers. Transforms — ScatterOutputs / GatherInputs / SliceOwned / AllReduceSum — say what to do per arg/output.

(3) Validate. trace_and_validate runs the model under a simulated 2-rank mesh, checks force-equivalence vs single-process, and — when a bare op needs an adapter — tells you which one to add:

from nvalchemi.distributed.validate import trace_and_validate
 
report = trace_and_validate(MyWrapper(model), world_size=2)
assert report.ok, report.next_action   # e.g. "add OpAdapter(..., ScatterOutputs())"

(4) Persist. MLIPSpec.save(path) / MLIPSpec.load(path) round-trip the spec so production runs skip re-discovery.

Full walkthroughs: examples/distributed/04_byo_pytorch_mpnn.py (autograd MPNN, no custom op) and 05_byo_graph_transformer.py (a custom-op kernel needing an OpAdapter). Deep reference: the Bring Your Own Model user guide.

4. Bring your own dynamics under DD

A custom {class}~nvalchemi.dynamics.base.BaseDynamics integrator (see nvalchemi-dynamics-implementation) runs under DomainParallel unchanged as long as every operation is per-atom. The one rule:

Under DD each rank's pre_update / post_update sees only its owned atoms. Per-atom math is correct as-is; any global reduction is not.

Safe locally (no change needed): position/velocity integration, per-atom Langevin friction + noise, applying forces. NVE and NVTLangevin are exact under DD for exactly this reason. Reading batch.energy is also fine — the forward already reduces + replicates it globally.

Must be globalized (a shard .sum() / .max() / .dot() is WRONG): total kinetic energy, temperature, degrees of freedom, FIRE dot-products (v·f, v·v, f·f), a global convergence test, barostat pressure.

For the shipped ensembles this is automatic: Nosé–Hoover, NPT, NPH, and FIRE declare their global quantities as intent and DomainParallel's coordinator reduces them — prefer reusing them. For a genuinely custom global scalar, compute it in a HookScope.GLOBAL hook, which sees the full gathered system on every rank:

from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.distributed import HookScope
 
class GlobalKineticEnergyHook:
    stage = DynamicsStage.BEFORE_STEP
    scope = HookScope.GLOBAL       # gather the FULL system to every rank first
    frequency = 1
    def __init__(self, integrator):
        self._integrator = integrator
    def __call__(self, ctx, stage):
        b = ctx.batch              # full system, identical on every rank
        self._integrator.global_ke = 0.5 * (
            b.atomic_masses * (b.velocities**2).sum(-1)
        ).sum()
 
# Register on the OUTER DomainParallel; the integrator reads self.global_ke.
dynamics = DomainParallel(dynamics=my_integrator, config=cfg, n_steps=n,
                          hooks=[GlobalKineticEnergyHook(my_integrator)])

Because every rank computed the same value from the same gathered system, the integrator stays in lockstep.

Don'ts:

• Don't hand-roll torch.distributed.all_reduce over the world group — route reductions through a shipped ensemble or a GLOBAL hook so they stay correct under a replicated layout (where every rank holds the full data).
• Don't read owned + ghost atoms inside pre_update / post_update — you only ever get owned atoms.
• Don't make a per-rank control-flow decision (like "converged") from a local scalar — DomainParallel already reduces convergence mesh-wide; a divergent local decision desyncs collectives.

Reference: key imports

See also

• User guide: Overview: Domain Decomposition, Bring Your Own Model, ShardTensor, and the Architecture & design deep dive.
• Examples: examples/distributed/03–07.
• Related skills: nvalchemi-model-wrapping (wrap a model), nvalchemi-dynamics-implementation (write an integrator), nvalchemi-dynamics-hooks (hooks), nvalchemi-dynamics-api (run/scale).