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. ```python from nvalchemi.distributed import ( DistributedManager, DomainConfig, DomainParallel, HookScope, ) ``` Launch with **torchrun** (or SLURM): DD is one process per GPU. ```bash 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()`. ```python import torch from nvalchemi.data import AtomicData, Batch from nvalchemi.distributed import ( DistributedManager, DomainConfig, DomainParallel, HookScope, ) from nvalchemi.dynamics import NVTLangevin, HostMemory from nvalchemi.dynamics.hooks import SnapshotHook from nvalchemi.dynamics.base import DynamicsStage from nvalchemi.hooks import NeighborListHook from nvalchemi.models.mace import MACEWrapper # 1. Bootstrap (reads RANK / WORLD_SIZE / LOCAL_RANK from torchrun). DistributedManager.initialize() dm = DistributedManager() mesh = dm.initialize_mesh(mesh_shape=(dm.world_size,), mesh_dim_names=("domain",)) device = torch.device(dm.device) # 2. Wrap the model — identical to single-process. wrapper = MACEWrapper.from_checkpoint("medium-mpa-0", device=device).eval() # 3. Build the inner integrator (its NeighborListHook is an INNER hook). integrator = NVTLangevin( model=wrapper, dt=1.0, temperature=300.0, friction=0.01, n_steps=200, hooks=[NeighborListHook(wrapper.model_config.neighbor_config, skin=0.5, stage=DynamicsStage.BEFORE_COMPUTE)], ) # 4. Trajectory snapshot: gather to rank 0 (an OUTER hook). snapshot = SnapshotHook(sink=HostMemory(capacity=201), frequency=10) snapshot.scope = HookScope.RANK_ZERO # 5. Wrap + run. cutoff = wrapper.cutoff makes the halo width exact. domain_cfg = DomainConfig(cutoff=float(wrapper.cutoff), skin=0.5, mesh=mesh) with DomainParallel(dynamics=integrator, config=domain_cfg, n_steps=200, hooks=[snapshot]) as dynamics: full_batch = build_full_system(device) if dm.rank == 0 else None owned = dynamics.partition(full_batch) # returns THIS rank's owned atoms dynamics.run(owned) DistributedManager.cleanup() ``` 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: ```python 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. | | **Halo** (default) | **Graph-partition** | |---|---|---| | Split by | space (spatial domains) | atom index (balanced node blocks) | | Each rank holds | owned atoms + a ghost halo | the **full geometry** (replicated) + owns a node slice | | Best for | models with a bounded cutoff you can hand a padded view (MACE, NequIP, LJ, Ewald, PME) | models that build their own neighbour list inside `forward` (UMA / eSCN-family) | | Comms | one halo exchange per step | per-layer feature all-gather + reduce-scatter | | Cell / migration | cell load-bearing; atoms migrate | cell is a plain input; no migration | ```python 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): | Preset | Strategy | |---|---| | `SPEC_MPNN_HALO` | halo — scatter-heavy MPNNs (MACE, NequIP, generic) **and any differentiable, autograd-force model** (incl. a BYO pure-PyTorch pair potential) | | `SPEC_LJ_HALO` | halo — the *shipped* Lennard-Jones wrapper's **opaque Warp kernel** (carries `OpAdapter`s) | | `SPEC_UMA_HALO` | halo — UMA/eSCN (local scatter) | | `SPEC_EWALD_HALO` / `SPEC_PME_HALO` | halo — long-range electrostatics | | `SPEC_DFTD3_HALO` | halo — DFTD3 dispersion | | `SPEC_MPNN_GP` | graph-partition — MPNNs under node partition | 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: ```python 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: ```python 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: ```python 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.** `FIRE` is deprecated; it is the only relaxation optimizer that works under `DomainParallel`, because `FIRE2` does not declare its reductions yet, so use `FIRE` there until `FIRE2` gains support. For a genuinely custom global scalar, compute it in a `HookScope.GLOBAL` hook, which sees the full gathered system on every rank: ```python 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,