[Draft] Add Distributed Inference and Domain Decomposition through Shard Tensor#122
Open
dallasfoster wants to merge 45 commits into
Open
[Draft] Add Distributed Inference and Domain Decomposition through Shard Tensor#122dallasfoster wants to merge 45 commits into
dallasfoster wants to merge 45 commits into
Conversation
The low-level domain-decomposition mechanism: the ShardTensor backend and construction, particle halo exchange and gather primitives, op transforms, placement, reshard, and per-system context.
The declarative MLIPSpec/DistributionSpec layer with op and method adapters and the graph padder, plus the three-layer public surface (mechanism ops, intent helpers, and the package entry point).
DistributedModel and DomainParallel, the spatial partitioner, the sharded batch, and domain configuration.
The compile bridge, compile-time neighbor refresh, and output consolidation that make the domain-decomposed forward compile-safe.
A trace-and-validate harness that checks a wrapper's distributed forward against a single-process reference and reports the adapters a model needs to run under domain decomposition.
In-tree copy of the upstream ShardTensor implementation the distributed core builds on, isolated under _core/_upstream and excluded from linting.
Wire domain decomposition through the model wrappers, dynamics, hooks, neighbor lists, and the data layer; add the distributed examples, the test suite, documentation, and the packaging and dependency updates that support the feature.
Config-driven benchmarks for the domain-decomposition path: a forward single-vs-multi force-equivalence runner and an end-to-end NVT runner, with per-model YAML configs for LJ, Ewald, PME, MACE, AIMNet2, and UMA.
Integrate upstream main (UMA fairchem-core wrapper, training subsystem + reporting hooks, toolkit-ops 0.4.0 pin, inflight-refill fix) into the domain-decomposition feature branch. Conflict resolution principles: - DD model wrappers (mace/aimnet2/ewald/pme/dftd3/lj/uma) keep all domain-parallel machinery (distribution_spec, node_energy_key, halo paths, CompilePolicy); upstream refinements folded in (shared autograd helpers in models/_utils, checkpoint-spec plumbing, training-aware autograd, tensile-positive stress = -virial/volume). - uma.py reconciles the two independent UMA wrappers into one: upstream's official fairchem API/structure as the base with the DD caps GraphPadder, distribution_spec, and distributed forward grafted in. - Hooks adopt the new slim HookContext + DynamicsContext/TrainContext split; NeighborListHook keeps its DD CUDA-graph capture path and torch.compiler.disable, with the capture gate keyed off upstream's neighbor-list method selector. - pyproject keeps the merged uv config (cu12/cu13 + mace/uma conflicts, ops 0.4.0-rc) plus DD tool config; uv.lock taken from main (dependency set identical). Repo-managed infra files updated only on main (.claude/skills, .github, Makefile, AGENTS.md) are intentionally left at the branch's versions. Validation: no conflict markers, all files parse, ruff clean, package imports (modulo the local toolkit-ops version skew). Full DD gate + NPT re-run pending on the 2-GPU box before push.
main's nvalchemi/distributed.py (DistributedManager + rank/world/device resolvers) was shadowed by our nvalchemi/distributed/ package, breaking 'from nvalchemi.distributed import DistributedManager' (training/hooks). Move it to nvalchemi/distributed/_runtime.py and re-export the five public symbols from the package; drop the shadowed module.
The merge fold of the former distributed.py into the package adds five public symbols (DistributedManager, PhysicsNeMoUninitialized...Warning, collective_device, resolve_global_rank, resolve_world_size) to nvalchemi.distributed.__all__; pin them in the stability test.
…move test_distributed_correctness imported nvalchemi.dynamics.hooks.neighbor_list; the shared hooks moved to nvalchemi.hooks (NVIDIA#65). Point at the new path.
Port slab_correction into the local DD electrostatics kernels so the PME and Ewald wrappers accept pbc/slab_correction on both single-GPU and under domain decomposition. New nvalchemi/models/_ops/electrostatics/slab.py: - slab_compute_partial_moments: registered custom op (single + batch) for the three global per-system slab moments (M, M2, Q). Owned-sliced + all-reduced by the DD layer (eager via OpAdapter, compile via current_dd_context), exactly mirroring the PME total_charge pattern. - compute_slab_correction_from_moments: per-atom energy/force/charge-grad/virial from the globally-reduced moments, in pure Torch matching the warp-kernel Yeh-Berkowitz + Ballenegger formulas (differentiable, compile-traceable). particle_mesh_ewald_from_total_charge and the Ewald wrapper forward add the slab term element-wise into their result tuples. Both wrappers register the slab moment ops in distribution_spec.custom_ops. Fix estimate_ewald_parameters call in test_ewald_staged_bindings (the helper used a stale rs_cutoff signature; the estimator now derives the splitting from positions + cell).
The stage-2 staged reciprocal bindings launch Warp kernels that require int32 batch indices; the batched test helper built them as int64.
Gates DistributedModel(PME/Ewald, slab_correction=True) energy + per-atom force equivalence to the single-GPU reference on a non-degenerate slab partition, proving the owned-only + all-reduced slab moments are halo-correct.
Bump the in-plane box / atom count so the 2-rank x partition has real remote atoms (halo does not cover the whole system), making the slab DD gate a strong correctness check rather than a trivial all-owned one.
Signed-off-by: Dallas Foster <dallasf@nvidia.com>
Signed-off-by: Dallas Foster <dallasf@nvidia.com>
Signed-off-by: Dallas Foster <dallasf@nvidia.com>
Contributor
|
Too many files changed for review. ( Bypass the limit by tagging |
dallasfoster
requested review from
atulcthakur,
laserkelvin and
zubatyuk
and removed request for
atulcthakur and
laserkelvin
A failed torch.cuda.graph capture in _NLGraphCache.capture() left the capture stream in an aborted stream-capture state, poisoning every later CUDA call with cudaErrorStreamCaptureInvalidated / "Offset increment outside graph capture" and cascading into unrelated CUDA tests. Wrap the capture region in try/except: on any failure, drop the partial graph, hard-synchronise the capture stream and device to clear the poisoned state, and latch _capture_disabled so the cache never retries (a retry would re-poison the context every step). The hook stays on the eager dispatch path, which is correct on its own — the captured graph is only a replay optimisation. capture() now returns a bool for the caller.
The cohesion refactor moved custom_ops under spec.distribution and changed distributed_setup to take a single ctx; the merged electrostatics wrappers use a fresh-tensor energy path (no _energies_buf alias). Update the stale tests to the current surface: - test_pme.py: TestPMEDistributionWiring uses spec.distribution.custom_ops (6 ops incl. slab moments), ctx-based distributed_setup, _dist_ctx / _n_global_atoms attrs; _energies_buffer test asserts the returned energy carries grad_fn while _energies_buf stays None. - test_ewald.py: same _energies_buf assertion update. - test_distributed_models.py: cueq spec tests read spec.distribution.custom_ops and resolve the spec without constructing a cueq model on CPU (skip cleanly without cuequivariance). - test_aimnet2.py: from_checkpoint loader calculator is train=False; the wrapper calculator is train=not compile_model. - test_distributed_correctness.py: build AtomicData cell as [1,3,3] and pbc as [1,3] per the current batched schema.
dallasfoster
requested review from
atulcthakur,
laserkelvin and
zubatyuk
and removed request for
laserkelvin
On warp/torch builds where warp's async mempool free is not capture-safe, a torch.cuda.CUDAGraph capture of warp kernels corrupts the CUDA context irrecoverably (later torch.randn calls raise "Offset increment outside graph capture"), cascading into unrelated GPU tests in the same process. Probe capture capability in a throwaway subprocess (so a failure can't poison the test process) and skip the four capture-requiring tests when unsupported. The eager NeighborListHook path is exercised by the rest of the suite; capture is an opt-in replay optimisation.
…CE training-flag tests AIMNet2Wrapper.from_checkpoint builds a throwaway loader calculator with train=False to extract the raw module; AIMNet2Calculator freezes (requires_grad_(False)) the shared module for train=False but never re-enables grad for train=True. The extracted module was therefore frozen, so a checkpoint loaded for training (train=not compile_model) came back non-fine-tunable. Re-enable requires_grad on the extracted module when train is requested, before handing it to the wrapper's own calculator. Tests: - test_aimnet2.py: adapt_output now strips per-atom padding via data.num_nodes, so the adapt_output tests pass the source batch instead of None; sized the raw per-atom tensors to num_nodes. - test_mace.py: the wrapper forward passes training=False to the inner MACE forward unconditionally (inference-only), so the training-flag tests assert [False]; fine-tuning still updates the trainable scale via autograd (assertion retained).
…aps) - test_train_mode_works_with_optimizer: xfail(strict) — the wrapper forward is inference-only (training=False), so forces have no grad_fn and a force loss cannot backprop. Energy-loss fine-tuning still works. - test_mace_cueq_dist_model_equivalence_2ranks: xfail(strict) — CONFIRMED real DD correctness gap in the cueq path (systematic ~1.8e-4 relative energy error that persists in fp64 and scales with N; plain MACE DD is machine-exact). Tolerance NOT loosened; xfail tracks the bug for fix.
…e system - model.model_card -> model.model_config (the wrapper attribute name). - Bump the argon system from 125 atoms (box 17A, degenerate for a 2-rank decomposition at ghost_width=cutoff8.5+skin1.0=9.5) to 343 atoms (box 23.8A > 2*9.5), so each rank owns atoms and the partition is a real decomposition.
…rd branch) The 'distributed' branch of the eager index_select / scatter_add dispatch guarded on isinstance(policy, PlainShard) and called policy.gather/scatter — but that branch is ONLY entered when the policy IS PlainShard (the classifier requires it), and PlainShard.gather/scatter intentionally raise NotImplementedError (storage-only). The correct gather-meta handlers (_distributed_index_select_handler / _distributed_scatter_add_handler) were therefore dead code, and every cross-rank sharded gather/scatter raised 'PlainShard is storage-only'. Route the distributed branch straight to the gather-meta handlers (the sharded path is driven by _gather_meta, not the storage policy).
…itions) Several distributed tests built a small cluster/chain inside a fixed oversized box (100 A cube, or n_per_side*spacing + 10 A), leaving atoms bunched in one corner. A 2-rank (or 4-rank) spatial bisection then put every atom on one side and assigned the other rank 0 owned atoms — a degenerate partition the framework correctly rejects, which masked the behaviour under test. Size the domain box to the atoms so the partition splits owned atoms onto every rank: - test_domain_parallel.py::_build_lj_cluster: box = n_per_side*spacing (drops the +10 A pad). Fixes test_nvt_langevin_lj_2ranks_end_to_end and the 4-rank NVE case. - validate/test_validate_cuda.py::_make_octane_chain_for_aimnet2: size cell_x to the chain extent (1024 atoms * 1.5 A ~ 1536 A >> 100 A box). - validate/test_helper_diagnosis.py::_make_octane_chain: size cell_x to the chain extent so both ranks own atoms.
A 1-D carbon chain (atoms only along x, at y=z=0) in a box with off-axis dims >= cutoff makes the partitioner build competing cells on y/z; the surface-area-minimising 2-rank grid can split along y or z, dropping every atom onto one rank and leaving the other with 0 owned (degenerate). Make the off-axis cell dims smaller than the cutoff (one cell each) so the split is forced onto x, the axis the chain actually spans, and both ranks own atoms. Fixes test_e2e_diagnostic_* and hardens the 1024-atom AIMNet2 validate sample.
… partition The 8-atom chain partitions trivially (each rank's halo covers the whole system, remote==0), so the unwrapped-mol_sum spoof produces no cross-rank divergence and the diagnostic has nothing to flag. A 24-atom chain (extent 36.5 A) gives each 2-rank domain > 2*ghost_width, so remote atoms exist and the per-system-reduction gap is detectable.
…ort imports Annotate intentional defensive try/except/pass (capture cleanup, fake-mode detect) and the capability-probe subprocess; sort imports in the distributed_correctness test.
MACE wrapper forward hardcoded training=False ('inference only'), so forces
carried no grad_fn and force/stress losses could not backprop. Thread
self.training so train mode retains the autograd graph while eval/inference/DD
keep the cheaper path. Un-xfail test_train_mode_works_with_optimizer.
cuequivariance conv fusion (enabled on CUDA by convert_e3nn_cueq) folds the InteractionBlock message pass — sender-gather + channel-wise TP + receiver-scatter — into one opaque kernel with the edge indices internal. That kernel gathers and scatters on rank-local indices and bypasses ShardTensor dispatch, so under domain decomposition it computes a purely local message: no halo gather-refresh of ghost senders, no reverse-exchange of ghost-receiver partials. The result is silently wrong on any non-degenerate partition (a near-degenerate cell masks it because the halo correction is a no-op there, and loose gate tolerances hid the residual). Fix it declaratively in the spec, keeping the model forward distribution-pure: a new ModuleForwardAdapter swaps a single module instance's forward for the DD scope and restores it on exit (MethodAdapter operates on a class; the fused forward is monkeypatched per instance). mace._cueq_conv_unfuse_adapters builds one per fused conv, declared on the cueq spec. The replacement reverts to the external node_feats[sender] gather + scatter_sum that plain MACE uses, so the halo read-refresh and scatter-correction fire (eager) / the static halo ops and message-passing refresh adapter fire (compile). Because the framework installs spec adapters only inside the DD scope, the replacement carries no DD branch and the single-process path keeps the fused kernel. The per-edge call into the fused descriptor reproduces the fused output to machine precision. The old _mace_cueq_spec hung scatter_outputs=[0] on fused_tensor_product, an op this path never calls; drop it — the cueq ops are all node/edge-local and the correction now lives in the unfuse adapter plus the external scatter's handler. Move both cueq gates to a non-degenerate 8x3x3 cell so the halo does real cross-rank work, remove the eager xfail, and update the docstrings. Validated on 2xA6000: fp64 machine-exact (eager dE=2.2e-11, compiled dE=3.6e-12); cueq and non-cueq are bit-identical in fp32 DD.
cuequivariance's EquivariantProductBasisBlock.forward derives the per-node element index for the symmetric contraction with ``index_attrs = torch.nonzero(node_attrs)[:, 1]``. Under the compiled domain-decomposition path the graph is padded to fixed-shape caps with inert dead atoms (Z=0, so an all-zero one-hot node_attrs row), so ``nonzero`` UNDERCOUNTS — ``index_attrs`` ends up with fewer rows than ``node_feats`` (n_padded) and the cueq ``uniform_1d`` kernel raises "batch dim mismatch" (caught in the op's fake impl while tracing on some torch.compile backends, at eager runtime on others). It is also data-dependent, which the compiled region cannot trace cleanly. Fix it in the spec: a MethodAdapter wraps the block forward and derives the index with ``node_attrs.argmax(dim=1)`` — one index per row, identical to ``nonzero[:, 1]`` for genuine one-hot rows, with dead rows mapping to element 0 (stripped from the owned-only output). Static shape, so it also sidesteps the data-dependent nonzero under compile. The non-cueq branch has no nonzero and delegates to the original; the adapter installs only inside the distributed scope, so single-process keeps the stock forward. Validated on 2xA6000 (cueq medium-mpa-0, non-degenerate world=2): eager and compiled DD match the single-process reference to ~1e-5 with no regression.
…peak The DD scaling sweep was measuring the wrong thing. run_sweep_one_size computed the single-rank full-system reference on EVERY rank before the multi-rank leg, so under domain decomposition every GPU built and forwarded the whole system: the reported multi-rank peak memory was the full-system peak (not the per-GPU shard), max-N was capped at the single-GPU OOM ceiling, and a reference OOM could take down the run before the DD forward executed. Add --run-reference (default OFF). The single-rank reference + equivalence check now run only when world_size==1 (where it IS the measurement) or when explicitly requested; otherwise the multi-rank leg is DD-only, so no rank holds the whole system and the reported peak/max-N reflect the shard. ``run_ref`` is identical on every rank, so the surrounding collectives stay symmetric. Also reset the CUDA peak AFTER warmup in time_step, so the reported peak is the steady-state per-step memory MD actually uses, not the one-time torch.compile / autotuning transient.
The cueq conv-fusion DD adapter unconditionally unfused the conv to the
external gather + scatter form so the halo dispatch handlers could fire.
Under torch.compile that materializes the per-edge message as a
saved-for-backward activation, roughly doubling compiled-DD peak memory
(25.5 -> 11.2 GB at 9k atoms, 2 ranks) and halving max-N vs single-GPU
compiled.
Make the conv adapter mode-dependent via compile_routing_active():
- eager DD: unfuse (external scatter; halo via dispatch handlers)
- compiled DD: stay fused; halo correctness comes from the refresh
adapter's scatter_to_owners on the block output
Both forms are machine-precision force-equivalent to the single-process
cueq reference. Compiled-DD peak now matches single-GPU compiled exactly
(11240 MiB @9k). Removes the NVALCHEMI_MACE_NO_UNFUSE debug env gate.
…icy dispatch Add a graph-parallel (edge-partition / replicate-nodes) domain-decomposition strategy alongside the spatial halo strategy, selectable per model via the distribution spec. Each rank owns a balanced index slice of atoms plus the edges into them; node features are all-gathered to a replicated tensor per message-passing layer (reduce-scatter adjoint), and the owned per-graph energy is summed locally then all-reduced. Complements halo: GP wins small/dense where halo ghost-fraction approaches one, halo wins large-N where GP's per-layer node all-gather dominates. - GraphParallelPolicy + IndexPartitioner + SPEC_MPNN_GP; gather_to_replicate collective with an exact all-reduce adjoint. - DistributedModel._call_graph_parallel: owned-target edge prep (global graph, keep owned-receiver edges, remap to owned-local), energy-only wrapper run, autograd forces from the owned partial, sharded consolidation. - Make StoragePolicy carry the per-strategy behavior (replicate/fold, build_topology, run_forward, partition_mode, serialization) so the framework dispatches without isinstance branches and a new strategy is bring-your-own via subclass + register_policy_kind. - system_sum works for any non-halo policy (mesh sourced from the context). - Gates: toy MPNN collective equivalence, and a bring-your-own toy wrapper run through DistributedModel; energy + owned forces match single-process at world 2/3.
…nto it
Add GraphReplicatePolicy, the node-replicate complement of the node-partition
GraphParallelPolicy: every rank holds the full node set and a sharded edge
slice, and each message-passing layer's partial per-receiver message sum is
recombined by a cross-rank all-reduce. It targets models whose forward can't be
expressed with the intent verbs — the model runs unchanged and the recombine is
injected purely by a declared adapter, so adopting it is a policy + adapter
declaration, not a forward rewrite, and it uses no ShardTensor.
- GraphReplicatePolicy (fold = all-reduce, replicate = identity, edges sharded);
registered for serialization.
- DistributedModel._call_graph_replicate: gather full nodes, shard edges, run
the wrapper energy-only, read the energy off each rank's owned node slice
(so the recombine's all-reduce adjoint sums distinct partials rather than the
replicated-energy gradient), forces = -dE/dx summed across ranks. A compile
branch (_gp_replicate_compiled_region) torch.compiles the energy-only forward;
the recombine traces inline as a mesh-static funcol all-reduce, so no per-step
routing tensors are threaded (simpler than the halo compile path).
- neighbor_refresh_adapters(always=True) fires the recombine in eager too.
- MACE: _mace_gp_replicate_spec (reuses the cueq/scripted passthrough + marshal
adapters, swaps the policy, declares the per-node energy, adds the conv
recombine), selected by NVALCHEMI_MACE_GP=1.
- Gates: toy opaque MPNN equivalence (world 2/3) and real MACE 2-GPU equivalence
across {non-cueq, cueq} x {eager, compile}, energy + forces vs single-process.
Add the GraphReplicate storage strategy for UMA: every rank holds the full node set plus a sharded edge slice; the model's internally-built graph is sharded by a declared _generate_graph slice adapter (otf_graph otherwise ignores injected edges), and the per-block edge->node aggregations (Edgewise, EdgeDegreeEmbedding) recombine across ranks via an all-reduce fold. Energy is read off each rank's owned node slice; forces/stress all-reduce in consolidation. The fold/reduction helpers stay policy-agnostic via the intent verbs (scatter_to_owners -> StoragePolicy.fold): identity under the refresh-only halo policy, all-reduce under graph-replicate. RefreshOnlyHaloPolicy distinguishes UMA's owned-complete halo aggregation (refresh-only) from a scatter-correction halo. The DistributionSpec lowers adapters onto third_party_helpers in __post_init__, so the GP spec passes ONLY the new slice adapter (the shared folds are already lowered via the halo spec) — re-passing them would double-install and apply the fold twice. ctx gains owned_offset (interior owned slice for node-replicate) and a mesh.get_local_rank() fallback for ctx.rank off the halo path. Validated: machine-precision force/energy equivalence vs single-GPU at world=2 on a rattled bcc-Fe cell; halo path unchanged.
Make UMA graph-replicate correct under torch.compile (fairchem turbo): - @torch._dynamo.disable the _generate_graph edge-slice adapter so the is_distributed gate + slice run eagerly per call, reading the live context, instead of baking the first-traced branch into the compiled graph. The heavy convs around it stay compiled (same pattern as the existing eager-refresh island). - Compiled UMA requires merge_mole=True: the per-forward MoLE state (mole_sizes / expert_mixing_coefficients) is mutated inside the compiled graph and lost, so the unmerged MoLE forward reads empty state. Merging folds experts into the weights once (compile-safe, algebraically exact); for graph-replicate the merge runs per rank over the full replicated node set, giving the correct global coefficients. - The equivalence test gets a UMA_GP_COMPILE gate (turbo-style settings) and a separate eager reference wrapper: sharing the GP wrapper would let fairchem compile the model on the reference call, before the DD adapters are installed, so torch.compile would capture the unpatched methods and the GP forward would silently run un-sharded. Validated: compiled GP force/energy machine-exact vs eager single-GPU at world=2 on a rattled bcc-Fe cell.
atulcthakur
requested changes
Jun 30, 2026
Comment on lines
+540
to
+541
| inv_cell = torch.linalg.inv(cell_3x3) | ||
| frac = batch.positions @ inv_cell.T |
Collaborator
There was a problem hiding this comment.
The codebase uses cart = frac @ cell, so this should multiply by inv_cell, not inv_cell.T. The current code wraps atoms incorrectly in skewed or triclinic cells. For non-orthogonal cells this changes positions even when they are already inside the box. Please add a skew-cell regression test after the fix.
| self._call_hooks(DynamicsStage.BEFORE_STEP, batch) | ||
|
|
||
| dyn = self._dynamics | ||
| dyn._ensure_state_initialized(batch) |
Collaborator
There was a problem hiding this comment.
NVE and Langevin can operate on local atom state, but Nosé–Hoover, NPH, and NPT depend on global kinetic energy, degrees of freedom, pressure, and shared cell state. This implementation computes those values independently from each rank’s shard, so ranks can evolve different thermostat/barostat states and cells. Please synchronize these quantities across the domain mesh or other less optimal option is to explicitly limit DomainParallel to NVE/Langevin.
Comment on lines
+476
to
+479
| for name in ("atomic_numbers", "atomic_masses", "velocities", "forces"): | ||
| val = getattr(batch, name, None) | ||
| if val is not None: | ||
| fields[name] = val |
Collaborator
There was a problem hiding this comment.
Initial partitioning supports arbitrary atom fields, but migration only moves this fixed list. Atomic charges and potentially any other custom fields can disappear after an atom crosses ranks; please migrate every existing per-atom field.
| self.step_count += 1 | ||
| dyn.step_count += 1 | ||
|
|
||
| converged = dyn._check_convergence(batch) |
Collaborator
There was a problem hiding this comment.
Convergence is checked only against the atoms on this rank. Different ranks can therefore disagree and enter different control flow; please make convergence a mesh-wide decision.
|
|
||
| if scope == HookScope.GLOBAL: | ||
| if dist.is_initialized() and getattr(batch, "energy", None) is not None: | ||
| dist.all_reduce(batch.energy, op=dist.ReduceOp.SUM) |
Collaborator
There was a problem hiding this comment.
HookScope.GLOBAL is documented to give hooks the complete batch, but this still passes each rank’s local shard. It also sums batch.energy, which is normally already globally reduced, so the reported energy can be multiplied by the number of ranks. Please gather the full batch over the domain mesh and avoid reducing already-global outputs again.
| else: | ||
| elem_diff = float((v64 - gv64).abs().max().item()) | ||
|
|
||
| ref_flat_sorted = v64.flatten().sort().values |
Collaborator
There was a problem hiding this comment.
Sorting the fully flattened output keeps only the set of numbers and loses which atom and vector component each number belonged to. A force array with values assigned to the wrong atoms or axes can therefore compare equal and make validation pass. Please match outputs using stable atom IDs and preserve the vector dimensions during comparison.
| self.restore() | ||
| raise | ||
|
|
||
| def restore(self) -> None: |
Collaborator
There was a problem hiding this comment.
Two distributed models can temporarily replace the same global helper. If the older model closes first, it restores the original function while the newer model is still active; when the newer model closes, it can restore the older replacement and leave it installed globally. This is reachable in persistent model pipelines that close models in forward order. Please track replacements per target using a shared stack or reference count.
| # ZarrData for incremental disk persistence. | ||
| n_frames_expected = (args.n_steps // args.snapshot_every) + 1 | ||
| trajectory_sink = HostMemory(capacity=n_frames_expected) | ||
| snapshot_hook = SnapshotHook( |
Collaborator
There was a problem hiding this comment.
Could we make this a rank-zero hook that gathers the full batch? SnapshotHook has no scope, so DomainParallel defaults it to local and writes only each rank’s owned atoms. Since the example later saves only rank 0’s sink, the resulting XYZ contains only rank 0’s shard on a real multi-rank run.
|
|
||
| # Reuse the SiO2 supercell builder from the benchmark suite — one | ||
| # canonical periodic test system across the distributed examples. | ||
| sys.path.insert(0, str(Path(__file__).parent)) |
Collaborator
There was a problem hiding this comment.
_benchmark_common.py only exists under benchmark/distributed, but this adds examples/distributed to the import path. The documented command therefore fails here with ModuleNotFoundError. Please move the shared helper into an importable module or import it from the correct location.
| dist.destroy_process_group() | ||
|
|
||
|
|
||
| if __name__ == "__main__": |
Collaborator
There was a problem hiding this comment.
Once the import above is fixed, the default Sphinx Gallery build will execute this example as __main__. Example 03 ignores NVALCHEMI_SPHINX_BUILD and calls init_process_group() without a torchrun environment, so the docs build will fail. Please add the same guard used by examples 01 and 02.
…haling overhead Add a node-partition graph-parallel strategy for UMA (NVALCHEMI_UMA_GP=partition): each rank runs the backbone on its owned atom block with a per-layer node-feature all-gather (reduce-scatter adjoint), via two eSCN adapters (_generate_graph owned-receiver slice + Edgewise all-gather) and LOCAL-scope energy/reference reductions. Forces come from the model's internal autograd over the full positions, summed across ranks with no /world (the feature all-gather's reduce-scatter backward routes each node's gradient to its owner exactly once). Force-equivalent to single-process at world=2. _call_graph_parallel now branches on forces_via_autograd: the framework-autograd owned-leaf path (toy MPNN, SPEC_MPNN_GP marked FRAMEWORK_FROM_GLOBAL_ENERGY) vs a new _graph_parallel_internal for models with opaque internal force heads. Cut the DistributedModel per-forward overhead that dominated graph-parallel time (272ms/forward at N=2000 -> ~2ms), bringing it level with fairchem native graph parallel: - ShardedBatch._build_batch_from_tensors uses model_construct instead of the validating AtomicData(...) constructor, whose atom_categories Enum-coercion calls repr on CUDA tensors (per-element host syncs). - _graph_parallel_internal mutates positions to the autograd leaf in place instead of rebuilding AtomicData/Batch; _call_graph_replicate likewise uses model_construct for its edge-sliced batch. An env-gated (NVALCHEMI_DD_PROFILE) section profiler is left in the graph-parallel forward for future perf work.
…traceable all-gather) The per-layer node-feature all-gather now has a fullgraph-traceable fixed-shape form so it fuses into the model's compiled forward instead of breaking out to an eager island. refresh_neighbors routes to fixed_gather_to_replicate (the funcol fixed gather-by-index with global_indices = arange(N), reduce-scatter-sum adjoint) when a node-partition compiled forward is active. The node-partition routing (owner_rank, local_index, cap) is index-based and static across MD steps, so it is published eagerly (set_gp_compile_routing) and read inside the compiled region as trace-time constants -- no graph-input threading, no recompiles at fixed N. The fixed gather is gated on torch.compiler.is_compiling(), so eager forwards keep the exact-size gather_to_replicate path. Compiled UMA node-partition is force-equivalent at world=2 (|dF| <= 8e-6) with zero steady-state recompiles across MD-like perturbed forwards.
…o DD + config-driven strategy selection Make the parallelization axis a first-class, hot-swappable Strategy that owns the whole vertical slice of strategy-dependent behavior, so Halo and graph-parallel share one seam instead of being re-decided across the stack. - S0: ParallelizationStrategy protocol + ShardState (nvalchemi/distributed/strategy.py); Halo/GraphPartition/GraphReplicate concretions wrapping the per-field StoragePolicy. - S1: migration + cell tracking move onto the strategy (plan/apply migration, on_cell_change); DomainParallel delegates, spatial hardcoding removed. - S2: DistributedModel forward branches relocate to strategy.run_forward (halo / graph-partition / graph-replicate); policy.run_forward inversion deleted. - S3: DynamicsDistributionCoordinator replaces GlobalThermoSync — NHC/NPT/NPH reductions route through strategy.reduce_system (node-replicate = identity, not an over-counting all_reduce); integrators declare intent via __dd_thermo_kind__ / __dd_replicated__. Includes the global-thermo consumer wiring + ops flag bindings. - S4a: config-driven strategy selection — StrategyKind + DomainConfig.strategy; distribution_spec becomes distribution_spec(strategy); NVALCHEMI_MACE_GP / NVALCHEMI_UMA_GP env-var gates removed. - S5: split ShardedBatch into a generic base + HaloShardState so each strategy's scatter builds its natural layout (halo owned+ghost; graph-parallel carries no halo baggage). Validated on 2×A6000 across the model×strategy matrix (eager equivalence, NHC/NPT dynamics, MACE/UMA graph-parallel) and in production execution modes (MACE cueq+compile, UMA turbo). Also refreshes the cueq custom-ops spec test to the current pass-through design and hardens the MACE NVT distributed example (docs-build guard + RANK_ZERO snapshot scope).
…caling harness
The forward + NVT benchmark harnesses selected only halo (they built
DomainConfig without a strategy); post-refactor that silently picks the halo
spec even for a graph-parallel layout. Wire config-driven selection:
- SystemConfig.strategy ("halo" | "graph_replicate" | "graph_partition") +
resolve_strategy() → (StrategyKind, partition_mode); both harnesses build
DomainConfig(strategy=...) and keep the ShardedBatch partition layout
consistent. Sweep a strategy with --set system.strategy=...
- Skip the upfront halo_exchange when there is no halo config (graph parallel),
fixing an AttributeError on the GP forward path.
Validated at world=2 on 2×A6000: MACE cueq halo and graph_replicate both run.
laserkelvin
requested changes
Jul 2, 2026
laserkelvin
left a comment
laserkelvin
left a comment
Collaborator
There was a problem hiding this comment.
Leaving some initial comments
| # ----- Process group setup ----- | ||
| # ``torchrun`` populates RANK / WORLD_SIZE / LOCAL_RANK; we just | ||
| # bind to the assigned device and init the process group. | ||
| dist.init_process_group(backend=args.backend) |
Collaborator
There was a problem hiding this comment.
Do you want to make this more DistributedManager oriented instead?
| # so single-GPU multi-rank correctness testing works end-to-end on | ||
| # Gloo. NCCL on a real multi-GPU cluster goes through the unmodified | ||
| # collective. | ||
| if args.backend == "gloo": |
Collaborator
There was a problem hiding this comment.
Should we even have gloo as a testable backend for this, and not just use nccl as the backend?
| warnings.simplefilter("ignore") | ||
| from nvalchemi.models.mace import MACEWrapper | ||
|
|
||
| dtype = torch.float32 |
Collaborator
There was a problem hiding this comment.
Maybe have these declared at the top, if you intended for people to be able to change the dtype
| ) | ||
|
|
||
| dynamics.close() | ||
| dist.destroy_process_group() |
Collaborator
There was a problem hiding this comment.
Not sure if you feel like it's valuable, but I think it'd be pretty cool to enclose the distributed semantics into either DomainConfig or DomainParallel, so that you could then have
dynamics = DomainParallel(...)
with dynamics:
# context manager handles initialization and teardown of distributed processes
dynamics.run(owned_batch)
Collaborator
There was a problem hiding this comment.
Should add a skip if NVALCHEMI_SPHINX_BUILD=1 env variable check so that the gallery doesn't try and run this
| else None | ||
| ) | ||
| # Energy: each rank holds its owned per-system partial → global SUM. | ||
| if "energy" in output and isinstance(output["energy"], torch.Tensor): |
Collaborator
There was a problem hiding this comment.
I'm wondering if it's worth fusing the energy and force reductions - the energy one is going to be latency bound anyway
| f = output["forces"] | ||
| if grp is not None: | ||
| f = f.clone() | ||
| dist.all_reduce(f, op=dist.ReduceOp.SUM, group=grp) |
Collaborator
There was a problem hiding this comment.
Instead of all_reduce, you could write it as a reduce-scatter so you don't need to communicate the full tensor
| raw["energy"] = total_energy | ||
|
|
||
| if compute_forces: | ||
| (forces_grad,) = torch.autograd.grad( |
Collaborator
There was a problem hiding this comment.
Didn't we have a helper function for this? Or is this just for transparency?
Same comment applies to the other BYO exampl
| :caption: The shipped presets cover the production model families. | ||
| from nvalchemi.distributed.spec import ( | ||
| SPEC_MPNN_HALO, # MACE, NequIP, Allegro, ORB | ||
| SPEC_AIMNET2_GATHER, # AIMNet2 (sharded, mol_sum reductions) |
Collaborator
There was a problem hiding this comment.
I don't think this was exported in that namespace
Collaborator
There was a problem hiding this comment.
I searched for ReplicatedMetadata symbol and couldn't find anything
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ALCHEMI Toolkit Pull Request
Description
Adds domain decomposition (DD) — multi-GPU spatial decomposition with halo
exchange — for nvalchemi MLIP inference and molecular dynamics, and integrates it
with current
main. A model author opts in purely via a declarativedistribution_spec;DomainParallelthen shards a system across ranks, exchangeshalo (ghost) atoms, and reduces owned-aware energies/forces/stresses so multi-rank
results match single-GPU to machine precision. Works eager and under
torch.compileacross MACE (scripted + cuEq), AIMNet2, UMA, DFTD3, PME, Ewald(incl. slab correction), and LJ, plus composed multi-model pipelines, and is
validated end-to-end through NVT and NPT MD.
Type of Change
Changes Made
nvalchemi/distributed/):DomainParalle spatial partitioner,ShardedBatch`, halo storage/exchange, owned-aware reductions,and a single-machine distribution validator.
MLIPSpec/DistributionSpecwithOpAdapter/MethodAdapter(owned-slice + all-reduce),CompilePolicy,GraphPadderfamily,and
node_energy_keyso the framework reduces per-node eeach declare a
distribution_spec; halo-only production path, hook-free.torch.compileunder DD: fixed-shape caps + dead-row padding + a compile-refreshgraph pass for zero steady-state recompiles in MD.
slab correction (
_ops/electrostatics/slab.py) thatreduced — verified DD-correct.
DistributedPipelineModel): direcand wired-charge (AIMNet2→PME) tiers, eager and compiled.
distributed/_core/_upstreamtounblock
torch.compile.main: UMA wrapper (feat(models): add UMA (fairchem-core) interatomic-potential wrapper #117), training subsystem + reporting hooks (Adding training functionalities to Toolkit #108),toolkit-ops 0.4.0 (Pin Toolkit-Ops version to 0.4.0 #120), inflight-refill fix (fix: preserve system IDs during inflight refill #113); tensile-positive stress
convention and the
nvalchemi.hooksnamespace adopted thTesting
Validated on a 2× RTX A6000 box (
nvalchemiops0.4.0,NCCL_P2P_DISABLE=1):test/distributed): 612 passed / 10 failed — all core force/energyequivalence (PME, Ewald, AIMNet2, MACE-scripted, pipeline, slab) passes; the 10
are peripheral (see Additional Notes).
test_slab_multigpu.py): world0-vs-world2ΔE = 0, max|ΔF| ≈ 3.5e-7 (fp32 floor) for both PME and Ewald.
multi-rank trajectory tracks single-GPU.
test/data(803) andtest/training(845) fully pass —the merge does not regress main.
make pytest) (see known test-debt below)make lint) on changed sourceChecklist
Additional Notes
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.