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().

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()