(distributed_design_overview)= # Distributed ML Potentials: Design Overview A guided tour of the distributed framework. Reads top-to-bottom as a 30-minute talk; each section is a slide group anchored by a figure and a code block. Cross-links into the deeper user-guide chapters where appropriate. | Section | Question it answers | |---|---| | 1. Motivation | Why does naïve domain decomposition break for MPNNs? | | 2. ShardTensor | What primitive lets distribution stay invisible to the model? | | 3. Specs | How does the framework know *what* to do at each op? | | 4. MACE end-to-end | What does "halo MPNN" actually look like? | | 5. UMA end-to-end | How do models that rebuild their own neighbour list scale? | | 6. Composition | How do MACE + Ewald run in the same pipeline? | | 7. Warp / Triton kernels | How do opaque kernels participate? | | 8. Validation + BYO | How does a new model author go from zero to production? | --- ## 1. Motivation: domain decomposition meets message passing ### 1.1 The starting point: classical DD works for short-range pair potentials Spatial decomposition is the standard way to scale molecular dynamics: each rank owns a region of the simulation cell, computes forces for its atoms, and exchanges a thin shell of "ghost" atoms with neighbors for pair interactions whose cutoff crosses the boundary. For a pair potential like Lennard-Jones, the math is local and the comms are cheap. ```{graphviz} :caption: Classical halo decomposition for a short-range pair potential. The global atom array is split into rank-owned contiguous slices; each rank materialises a thin shell of remote atoms within ``cutoff`` of its boundary (dashed). Pairs that cross the boundary are evaluated locally on either rank, no message is in flight at force-eval time. :align: center digraph halo_classical { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; global_view [label=<

global atom array (16 atoms)
0	1	2	3	4	5	6	7	8	9	10	11	12	13	14	15
rank 0 owns								rank 1 owns

>]; rank0_view [label=<

rank 0's local view: 8 owned + 4 shell rows
0	1	2	3	4	5	6	7	8	9	10	11
my owned								shell (read-only)

>]; rank1_view [label=<

rank 1's local view: 4 shell + 8 owned rows
4	5	6	7	8	9	10	11	12	13	14	15
shell (read-only)				my owned

>]; global_view -> rank0_view [label=<slice + add shell>]; global_view -> rank1_view [label=<slice + add shell>]; } ``` ```{code-block} python :caption: The mental model: per-rank locality holds because pair interactions decay with distance. # Single-process: O(N) atoms, O(N²) pairs (with cutoff: O(N)) for atom in batch: for neighbour in atom.within(cutoff): accumulate_force(atom, neighbour) # Halo distributed: each rank does the same loop on its # (owned + halo) atoms, only writing forces for owned. ``` ### 1.2 What breaks for message-passing potentials A scatter-heavy MPNN like MACE doesn't have one cutoff and one pair sum. It has L message-passing layers, each scattering edge features into per-atom features. After L layers, every atom's representation depends on the L-hop neighbourhood — even atoms whose owned-rank is *not* the receiver. Two failure modes: ```{graphviz} :caption: Halo width grows with depth. Each layer's update at an owned atom reads its 1-hop neighbours, so an owned atom's value at layer ℓ depends on the ℓ-hop neighbourhood. To compute correct values for every owned atom, the shell must reach ℓ × cutoff out from the rank boundary at layer ℓ. :align: center digraph halo_growth { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; layer1 [label=<

after layer 1: shell of width 1·cutoff is enough
0	1	2	3	4	5	6	7	8	9
my owned (correct features)								1-hop shell

>]; layer2 [label=<

after layer 2: layer-1 features for shell atoms 8..9 must be correct, so shell must extend to atoms 10..11 too
0	1	2	3	4	5	6	7	8	9	10	11
my owned								1-hop shell		2-hop shell

>]; layerL [label=<

after layer L: shell width = L · cutoff (gets expensive fast)
0	1	2	3	4	5	6	7	8	9	10	11	12	…
my owned								shell grows with L

>]; layer1 -> layer2 [label=<go to next layer>]; layer2 -> layerL [label=<…>]; } ``` The framework dodges this growth by *refreshing* the shell between layers (cheap exchange of one row's worth of data per shell atom) rather than expanding it. After layer 1 finishes, the shell rows hold correct *layer-1* features; layer 2's 1-hop reach into them is still correct. The two natural extremes that don't work: | Strategy | What happens | |---|---| | **All-gather every layer.** | Comms are O(N · L). Every rank sees the global tensor every layer. Parity with single-process but no scaling. | | **Strict-local.** Drop edges crossing the rank boundary. | No comms, but every owned atom near the boundary has missing neighbours. Forces on boundary atoms are wrong; energy is wrong; trajectories diverge. | ### 1.3 Beyond MPNNs: the long-range tail Modern ML potentials don't stop at MPNN. Several patterns make naïve locality even harder: | Pattern | Locality breaks because | |---|---| | **Charge equilibration / electrostatic embedding (AIMNet2)** | A per-system reduction in every layer. ``mol_sum(per-atom)`` is global. | | **Reciprocal-space methods (Ewald, PME)** | The structure factor / charge mesh is global. FFT is global. | | **Attention-based potentials** | All-pairs interactions, full softmax. | | **Graph-rebuilding models (UMA / eSCN)** | The model constructs its own neighbor list inside ``forward`` — there's no place outside the model to define a halo. | | **Stress via strain trick** | Differentiates a *replicated* per-graph energy through per-atom positions. | A halo-only world can't host these without per-pattern surgery. We need a primitive that lets each model declare *its* locality contract. ### 1.4 Where this lands The framework supports three storage strategies. Each is a different answer to the "where does each rank's per-atom tensor live, and what's the row layout?" question. ```{graphviz} :caption: Three ways to lay out a per-atom tensor across two ranks. Solid blocks are owned; dotted blocks are remote rows. **Halo storage** materialises a thin shell of remote owners' rows on each rank (read-only mirrors). **Sharded storage** stores only owned; cross-rank reads route over the wire. **Replicated storage** stores the full tensor on every rank and partitions logically — the layout used by the **graph-partition** strategy (§5), for models that rebuild their own NL inside ``forward`` and need to see every position. :align: center digraph storage_modes { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; title_halo [label=<HALO STORAGE
(MACE / NequIP / LJ / Ewald / PME)>]; title_sharded [label=<SHARDED STORAGE
(AIMNet2)>]; title_replicated [label=<REPLICATED STORAGE
(UMA / eSCN-family)>]; halo_r0 [label=<

	rank 0's owned		rank 1's shell
rank 0's local view: 8 owned + 4 halo rows

>]; halo_r1 [label=<

rank 0's shell		rank 1's owned
rank 1's local view: 4 halo + 8 owned rows

>]; sharded_r0 [label=<

	rank 0's owned
rank 0's local view: 8 owned rows

>]; sharded_r1 [label=<

	rank 1's owned
rank 1's local view: 8 owned rows

>]; repl_r0 [label=<

rank 0's partition		rank 1's partition
rank 0's local view: full 16 rows; node_partition = [0..7]

>]; repl_r1 [label=<

rank 0's partition		rank 1's partition
rank 1's local view: full 16 rows; node_partition = [8..15]

>]; title_halo -> halo_r0 [style=invis]; halo_r0 -> halo_r1 [label=<refresh shell from
each other's owners> dir=both color="#666"]; title_sharded -> sharded_r0 [style=invis]; sharded_r0 -> sharded_r1 [label=<cross-rank read or scatter
routes over the wire
(no shell stored)> dir=both color="#666" style=dashed]; title_replicated -> repl_r0 [style=invis]; repl_r0 -> repl_r1 [label=<per-MP-layer all_gather
of partition rows
(every rank sees the
full feature tensor)> dir=both color="#666"]; } ``` Quick reference for picking a strategy: | Strategy | Per-atom tensor row layout | When the model needs it | Per-step comm | |---|---|---|---| | **Halo** | ``[owned │ halo]`` — owned rows are unique, halo rows mirror remote owners. | Local-receptive-field MPNN where the cutoff fits in one halo width. The model can be handed an opaque ``(n_owned + n_halo, *F)`` view and produce correct outputs without knowing the partition. | One halo exchange per step (refresh shell rows). | | **Sharded** | ``(n_owned, *F)`` — each rank stores only its owned rows. | Per-system reductions inside every layer (charge equilibration; ``mol_sum``). The cost of a per-layer halo refresh is dominated by the global reduction anyway, so saving on memory wins. | One ``all_to_all_v`` per cross-rank read or scatter (already needed for ``mol_sum``). | | **Replicated** (graph-partition) | Full ``(n_global, *F)`` on every rank; partition is logical (a ``node_partition`` index slice). | Models that build their own neighbor list inside ``forward`` and can't be handed a pre-padded view (``UMA``'s ``_generate_graph``, ``eSCN``-family). Memory is O(n_global) per rank, so capped by single-GPU budget. | Per-MP-layer feature ``all_gather`` (autograd-aware) plus per-system ``all_reduce`` for energy / forces / stress. | The storage strategy is just the start. Within each strategy we still need to pick scatter rules, gather rules, and per-output reductions. Encoding those choices is the job of the *spec* (§3). --- ## 2. ShardTensor: a partition-aware Tensor subclass ### 2.1 The primitive `ShardTensor` is a `torch.Tensor` subclass that carries metadata about the partition, plus a registry of dispatch handlers that intercept specific torch ops on it via `__torch_function__`. The same model code runs single-process and distributed; the runtime decides what to do based on the input's metadata. ```{graphviz} :caption: A ShardTensor is a regular tensor block with a small bag of metadata attached. The data lives in the same buffer as a plain tensor (zero-copy). The metadata describes what the rows mean and which rules govern operations on them. :align: center digraph shardtensor_anatomy { rankdir=LR; node [shape=plaintext fontname="Helvetica"]; tensor [label=<

tensor data    (this rank's view, shape (12, 3))
	owned		shell
rows 0..7 are this rank's; rows 8..11 mirror neighbour rank's owned

>]; metadata [label=<

metadata bag
how many rows are mine	8
how many rows total (mine + shell)	12
which rank am I	0
where do shell rows come from	rows 8..11 ↤ rank 1's owned
how many systems (for per-graph reductions)	1
partition rules (see §3)	storage = halo scatter rule = halo_correction gather rule = halo_read per-system reductions = on

>]; tensor -> metadata [style=dashed label=<carries...>]; } ``` ```{code-block} python :caption: Construction. ``wrap`` is the canonical entry point — the framework calls it when promoting padded positions, and users call it for their own per-atom fields. from nvalchemi.distributed.ops import ShardTensor t = ShardTensor.wrap( halo_padded_positions, # (n_padded, 3) plain tensor spec=SPEC_MPNN_HALO, # tells handlers HOW to dispatch meta=halo_meta, # owned/padded counts, routing config=halo_config, # mesh + process group n_systems=1, # for per-system reductions ) ``` ### 2.2 Dispatch via `__torch_function__` When a torch op is called with at least one ShardTensor argument, PyTorch invokes `ShardTensor.__torch_function__`. We walk a small registry of `(op, predicate, handler)` tuples; the first matching predicate wins. No match → `super().__torch_function__` falls back to plain torch behaviour with metadata propagation. ```{graphviz} :caption: Dispatch flowchart for a single torch op. Predicates inspect the inputs' shapes + spec; the matching handler runs cross-rank logic and re-promotes outputs. :align: center digraph dispatch { rankdir=TB; node [shape=box style="rounded,filled" fontname="Helvetica"]; Op [label="op(shard_tensor, ...)" fillcolor="#dce6f1" fontcolor="#111111"]; Pred [label="any registered handler\npredicate matches?" fillcolor="#fff2cc" fontcolor="#111111" shape=diamond]; Handler [label="handler runs:\n• unwrap inputs\n• run cross-rank logic\n• promote outputs" fillcolor="#cfe2cf" fontcolor="#111111"]; Fallback [label="super().__torch_function__\n(plain torch.Tensor path)" fillcolor="#f3f3f3" fontcolor="#111111"]; Done [label="return result\n(metadata propagated)" fillcolor="#82b366" fontcolor="#111111"]; Op -> Pred; Pred -> Handler [label="yes"]; Pred -> Fallback [label="no"]; Handler -> Done; Fallback -> Done; } ``` ```{code-block} python :caption: A user-facing example. The wrapper code is identical to single-process; the dispatch machinery reads the partition off the tensor. # Single-process or distributed — same line: total_energy = torch.zeros(n_graphs, ...) total_energy = total_energy.scatter_add_(0, batch_idx, atomic_energies) # ^ # If atomic_energies is a ShardTensor with system_reductions=True # and accumulator shape == n_systems, the dispatch routes through # ``per_system_reduce``: local scatter + cross-rank all_reduce. # Otherwise: plain in-place scatter_add_. ``` ### 2.3 Why "almost transparent" There's a footgun the framework has to surface explicitly. PyTorch's in-place ops (`t.scatter_add_(...)`) return `self`; the standard idiom drops the return. Under cross-rank halo correction the handler *can't* preserve in-place semantics — it returns a fresh tensor. Wrapper authors must rebind the return: ```{code-block} python :caption: The one rule wrapper authors need to internalise. Single-process: rebind is a no-op. Distributed: mandatory. # WRONG (single-process: works; distributed: silently produces zeros) G.scatter_add_(0, receivers, edge_feats) # RIGHT G = G.scatter_add_(0, receivers, edge_feats) ``` The validator's worker-error translator detects this exact pattern and surfaces it as a one-line fix. See {ref}`distributed_byo_guide` for the full catalogue of dispatch handlers and their predicates. --- ## 3. Specs: declaring what to do at each op ShardTensor knows *how* the data is partitioned. The {py:class}`~nvalchemi.distributed.spec.MLIPSpec` tells it *what to do* at each op. The split is deliberate: ShardTensor stays chemistry-free; specs encode model-specific reduction rules. The spec is a small structure with three knobs. Each knob has a visual interpretation on the tensor — that's what the rest of this section walks through. | Knob | Choices | What gets visualised | |---|---|---| | **scatter rule** | `halo_correction` / `local` / `distributed` | how an `t.scatter_add_(...)` call moves data | | **gather rule** | `halo_read` / `local` / `distributed` | how an `index_select` call sees data | | **per-system reductions** | `on` / `off` | how a per-graph energy `scatter_add` becomes a global sum | Plus per-op transforms (§7) and per-output classifications (§8). ### 3.1 The scatter rule: where do partial messages go? A scatter is "for each edge, write a contribution into the receiver's row." When the receiver might be a halo row (a mirrored copy of a remote rank's owned), the scatter rule decides whether (and how) to account for that. ```{graphviz} :caption: ``scatter = "halo_correction"`` — the canonical MPNN pattern. Each rank scatters its messages into both owned and shell rows. The shell partials get sent back to their owners and accumulated. Then owners' values are pushed back out to refresh shell copies for the next layer. :align: center digraph scatter_halo { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; edge [color="#666"]; s0 [label=<

step 1: each rank does a local scatter — partials land in BOTH owned and shell rows
rank 0 owned filled with this rank's edges	rank 1's shell on rank 0 partials destined for rank 1

>]; s1 [label=<

step 2: send shell partials back to owners; owners accumulate
rank 0 owned unchanged	rank 1 owned receives contributions from rank 0

>]; s2 [label=<

step 3: owners broadcast their final values back into the shell so the next layer sees them
rank 0 owned	shell now mirrors rank 1's final owned

>]; s0 -> s1 [label=<send back
(reverse)>]; s1 -> s2 [label=<refresh
(forward)>]; } ``` ```{graphviz} :caption: ``scatter = "local"`` — pure per-rank scatter, no cross-rank exchange. Used when the accumulator is per-system (small) and the per-system all-reduce in step 2 of the next rule handles cross-rank correctness; or for a halo-unaware backbone whose edges already cover the global graph. :align: center digraph scatter_local { rankdir=LR; node [shape=plaintext fontname="Helvetica"]; a0 [label=<

rank 0 owned	shell
scatter_add_ writes locally; nothing crosses ranks

>]; a1 [label=<

shell	rank 1 owned
same: pure local

>]; a0 -> a1 [style=invis]; } ``` ### 3.2 The gather rule: how does an index_select see data? A gather is "for each input row index, fetch that row's data." When the index falls in the shell region (i.e. asks for a remote rank's row), the gather rule decides whether to serve from the local mirror or to route a request to the owner. ```{graphviz} :caption: ``gather = "halo_read"`` — the index ``9`` (in this rank's shell) is served from the local shell copy. No cross-rank traffic at gather time. Stays cheap because the shell is refreshed by the previous scatter's step 3. :align: center digraph gather_halo { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; tensor [label=<

0	…	7	8	9	10	11
rank 0's view: indices 0..7 are owned, 8..11 are shell

>]; request [label=<request: index_select [9]>]; served [label=<served from local shell row 9
(no cross-rank message)>]; request -> tensor [label=<read>]; tensor -> served [label=<return>]; } ``` ### 3.3 Per-system reductions: per-rank scatter + cross-rank sum The most common reduction in MLIPs is `total_energy.scatter_add_(0, batch_idx, atomic_energies)` — collapsing per-atom energies into a per-graph total. Under partitioning, no rank has all the atoms, so the local scatter is a partial. `per_system_reduce` does the local scatter, then sums the partials across ranks. ```{graphviz} :caption: ``per_system_reduce`` — one primitive that combines a local per-system scatter with a cross-rank sum. The output is replicated globally on every rank, so any rank can read the final per-graph value. :align: center digraph per_system { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; inputs [label=<

per-rank atomic energies (sliced to owned only)
rank 0:	e₀ e₁ e₂ e₃ e₄ e₅ e₆ e₇
rank 1:	e₈ e₉ e₁₀ e₁₁ e₁₂ e₁₃ e₁₄ e₁₅

>]; locals [label=<

step 1: local scatter into per-graph slot
rank 0:	Σ₀..₇
rank 1:	Σ₈..₁₅

>]; global [label=<

step 2: all_reduce sum — every rank holds the global total
rank 0:	E_global
rank 1:	E_global

>]; inputs -> locals [label=<scatter_add_ on owned slice>]; locals -> global [label=<all_reduce(SUM)>]; } ``` ### 3.4 The complete spec The decisions above all live on a small data structure that the wrapper attaches via `distribution_spec`: | Field | Purpose | |---|---| | `distribution.policy` | Storage layout: halo / sharded / local. Each carries its own scatter and gather rules. | | `distribution.custom_ops` | Per-op declarations for opaque kernels that bypass `__torch_function__` (Warp, Triton). See §7. | | `output_kinds` | One of `PER_NODE`, `PER_GRAPH`, `GLOBAL`, `UNKNOWN` per output. Drives final consolidation. | | `owned_only_outputs` | Per-atom outputs that are already globally correct on each rank (e.g. PME reciprocal forces) — skip the back-exchange. | | `all_reduce_outputs` | Per-rank-partial outputs that need a final SUM across ranks (e.g. strain-trick stress). | ```{code-block} python :caption: The shipped presets cover the production model families. A wrapper author either picks one or composes a new ``MLIPSpec`` directly. from nvalchemi.distributed.spec import ( SPEC_MPNN_HALO, # MACE, NequIP, Allegro, ORB (spatial halo) SPEC_MPNN_GP, # MPNN node-partition graph-parallel SPEC_UMA_HALO, # UMA / eSCN (spatial halo) SPEC_LJ_HALO, # Lennard-Jones (Warp pair kernel) SPEC_EWALD_HALO, # Ewald (real + reciprocal stages) SPEC_PME_HALO, # PME (charge spread + FFT mesh) SPEC_DFTD3_HALO, # DFT-D3 dispersion ) ``` Authoring a spec for a new model is the topic of §7 + §8. For now, note that **every spec parameterises the same dispatch machinery** — the registry, the predicates, the handlers. The spec is the single declaration point. --- ## 4. MACE end-to-end: what halo MPNN looks like ### 4.1 The forward pass, three steps ```{graphviz} :caption: One MACE message-passing layer under halo storage, viewed as tensor states (rank 0 of 2). The features tensor enters the layer as ``(n_padded, F)`` with both owned and shell rows populated. Edge messages scatter into the receivers' rows, leaving partial accumulations in shell positions destined for rank 1. The framework's scatter rule sends those partials back, then refreshes the shell so the next layer reads correct values. :align: center digraph mace_layer { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; fpre [label=<

features entering layer ℓ: shape (n_padded, F)
my owned (8 rows of features)	shell (4 rows mirroring rank 1)

>]; scatter [label=<

compute messages, scatter into receivers: writes land in BOTH owned and shell rows
owned: full sum of incoming edges where receiver is mine	shell: partial sum of edges where receiver lives on rank 1

>]; after_back [label=<

shell partials sent back to owners; rank 1's owners now hold the cross-rank contributions too
owned: unchanged	shell: emptied (rank 1 has the data now)

>]; fpost [label=<

features leaving layer ℓ: shape (n_padded, F), shell refreshed for layer ℓ+1's gather
owned: layer-(ℓ+1) features	shell: rank 1's layer-(ℓ+1) features

>]; fpre -> scatter [label=<edge messages>]; scatter -> after_back [label=<send back
(reverse exchange)>]; after_back -> fpost [label=<refresh shell
(forward exchange)>]; } ``` ```{code-block} python :caption: What the wrapper actually writes. The halo-correction is implicit in ShardTensor dispatch — the wrapper has zero distribution code. # Inside MACE InteractionBlock — a typical scatter pattern: node_feats = node_feats.zero_() node_feats = node_feats.scatter_add_(0, receivers, edge_messages) # ^ ^ # rebind handles if edge_messages is a ShardTensor, # distributed return the dispatch handler does # halo_reverse + halo_forward ``` ### 4.2 The full forward (one slide) ```{graphviz} :caption: A complete forward pass viewed as tensor states (rank 0 of 2). Positions enter as a halo-padded ShardTensor; L message-passing layers each apply the §4.1 pattern; the final atomic energies get sliced and reduced into a globally replicated total energy. Forces fall out of an autograd backward through positions; the framework routes shell gradients back to their owners. :align: center digraph mace_full { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; pos [label=<

positions: ShardTensor (n_padded, 3)
my owned (8 atoms)	shell (4 atoms, mirrored)

>]; feats [label=<

features after L MP layers: ShardTensor (n_padded, F)
my owned (correct)	shell (mirrored, layer-L)

>]; atE [label=<

atomic energies: head(features) → (n_padded,)
my owned: 8 per-atom energies	shell: 4 mirrors (will be dropped)

>]; pers [label=<

per-system reduce: drop shell, scatter into per-graph slot, all_reduce
local sum on rank 0:	Σ owned
after all_reduce:	E_global (replicated on every rank)

>]; bwd [label=<

autograd.grad(E_global, positions): produces ∂E/∂x on (n_padded, 3)
owned: this rank's contribution	shell: this rank's partial gradient at others' atoms

>]; forces [label=<

forces after consolidation: shell gradients sent back to owners

my owned forces (complete: own + cross-rank pieces)

>]; pos -> feats [label=<L message-passing layers
(each: scatter + send-back + refresh)>]; feats -> atE [label=<read-out head>]; atE -> pers [label=<scatter_add into (n_graphs,)
(per-system reduce)>]; pers -> bwd [label=<.backward()>]; bwd -> forces [label=<framework routes
shell gradients to owners>]; } ``` ### 4.3 Why this is short to write ```{code-block} python :caption: A halo-MPNN wrapper has zero distribution-aware code. The framework promotes per-atom inputs to ShardTensor before calling the wrapper; ``__torch_function__`` propagates the partition through the wrapper's ops; consolidation handles the final per-output reduction. class MACEWrapper(nn.Module, BaseModelMixin): @property def distribution_spec(self): return SPEC_MPNN_HALO # halo correction + halo read + per-system reductions def adapt_input(self, data, **kwargs): # Drop neighbour-list sentinel rows. Single-process: drops the # genuine padding rows the NL builder emits. Distributed: also # drops halo-receiver rows the framework rewrote to the same # sentinel value at NL-build time. One line, both regimes. n_atoms = data.positions.shape[0] edge_index = data.neighbor_list.long().T valid = (edge_index[0] < n_atoms) & (edge_index[1] < n_atoms) return { "positions": data.positions, # already a ShardTensor under DD "edge_index": edge_index[:, valid], "node_attrs": self._node_attrs(data), "shifts": ..., } def forward(self, data): return self.model(**self.adapt_input(data)) ``` The framework handles every cross-rank thing: halo build, NL filter, per-layer scatter/refresh, per-system reduction, force consolidation, strain-trick stress (with the inner virial pass routed correctly). --- ## 5. UMA end-to-end: node-partition graph parallel for graph-rebuilding models Some models can't be handed a halo-padded view because they build their own neighbor list inside ``forward``. UMA / eSCN-family models take ``positions`` and emit ``edge_index`` via their internal ``radius_pbc`` kernel — there's no pre-forward seam to attach a halo to. The **graph-partition strategy** ({py:class}`~nvalchemi.distributed.strategy.GraphPartitionStrategy`, selected with ``DomainConfig(strategy=StrategyKind.GRAPH_PARTITION)``) answers that: every rank holds the full positions tensor, the model's NL builder runs on the global geometry, then a balanced *node partition* — a contiguous slice of ``arange(n_global)`` — assigns each rank a distinct block of owned atoms. Each rank runs the backbone on its owned block. A per-MP-layer feature ``all_gather`` reconstructs the full node set the convolution needs, and a reduce-scatter adjoint on the backward routes each owned atom's cross-rank gradient back to its owner. Per-system energy and stress sum the owned slices with an ``all_reduce``; forces come from fairchem's own autograd (the ``MODEL_INTERNAL`` force strategy). Unlike the spatial halo the partition is geometry-free: the cell is an ordinary model input, atoms never migrate, and only the edge count drifts under MD (so compiled runs cap edges, not atoms). We don't reimplement fairchem's message passing. UMA's ``distribution_spec(StrategyKind.GRAPH_PARTITION)`` returns a spec whose ``policy`` is {py:class}`~nvalchemi.distributed.spec.GraphParallelPolicy` and whose adapters are a handful of ``MethodAdapter`` swaps that make the backbone owned-block-aware — leaving ``fairchem.core.common.gp_utils`` untouched (an earlier design redirected gp_utils via a thread-local metadata object; the node partition owns its own gather/reduce, so it no longer needs to): ```{code-block} python :caption: Pseudocode — the node-partition adapter set (built lazily inside UMAWrapper.distribution_spec). eSCNMDBackbone / Edgewise / ElementReferences are the real imported fairchem classes; each MethodAdapter uses the class-form ``MethodAdapter(RealClass, "method_name", replacement)`` (equivalently the keyword form ``MethodAdapter(module_path=..., class_name=..., method_name=..., replacement=...)``) to swap one eSCN method for a distribution-aware variant. partition_helpers = ( # replicate the full geometry, build the graph, keep this rank's owned nodes MethodAdapter(eSCNMDBackbone, "_generate_graph", _distributed_partition_graph), # all-gather owned node features to the full set for the edgewise conv MethodAdapter(Edgewise, "forward", _distributed_edgewise_gather), # undo element-reference offsets on the owned slice only MethodAdapter(ElementReferences, "undo_refs", _distributed_undo_refs), ) ``` The wrapper itself stays distribution-agnostic — its ``forward`` is the ordinary fairchem call. ``distribution_spec`` only picks the layout: ```{code-block} python :caption: distribution_spec returns the spatial-halo preset by default, or a node-partition GP spec (GraphParallelPolicy + the adapters above, built and memoised lazily) when the config selects GRAPH_PARTITION. class UMAWrapper(nn.Module, BaseModelMixin): def distribution_spec(self, strategy=None): if strategy == StrategyKind.GRAPH_PARTITION: # MLIPSpec(distribution=DistributionSpec(policy=GraphParallelPolicy()), # adapters=partition_helpers, outputs=...) — memoised here ... return SPEC_UMA_HALO # default: spatial halo def forward(self, data): return self.predict_unit(data) # fairchem does the rest ``` What the framework adds on top: * {py:class}`~nvalchemi.distributed.strategy.GraphPartitionStrategy` records the balanced node partition (``arange(n_global).tensor_split(W)[rank]``), replicates positions to every rank (no halo padding), runs the wrapper on the owned block, and consolidates the outputs — ``all_reduce`` for per-system energy / stress, reduce-scatter for owned force rows. * The partition is fixed for the run: no cell tracking, no migration. Only the per-rank edge count moves, so a compiled forward caps edges. Even though every rank holds the full positions tensor, per-rank MP-layer activations only span ``n_owned`` rows, so peak memory under 2 ranks is consistently 0.55–0.90× single-rank memory (better at larger N, where the activations dominate the peak over the replicated positions). The compute speedup is more modest (~1.20× forward, ~1.55× NVT at 2 ranks) because the per-MP-layer ``all_gather`` is in the critical path. --- ## 6. Composing models: pipelines that mix strategies Real workflows compose models. Energy = MACE (short-range MPNN) + Ewald (long-range electrostatics). Different sub-models can want different storage strategies and different specs. ```{graphviz} :caption: A two-block pipeline: MACE (short-range MPNN) + Ewald (long-range electrostatics). Halo construction happens once on the input; both blocks read the same padded tensor. Each block produces a globally replicated total energy and per-rank-owned forces; the pipeline sums them. :align: center digraph pipeline { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; input [label=<

input: per-rank owned atoms (8 each, ShardedBatch)
rank 0 owned	rank 1 owned

>]; padded [label=<

padded batch: built once, reused by both blocks
rank 0: owned + shell	rank 1: shell + owned

>]; mace [label=<

MACE block: short-range MPNN spec = SPEC_MPNN_HALO
energy:	E_MACE (replicated)
forces:	F_MACE (per-rank owned)

>]; ewald [label=<

Ewald block: long-range electrostatics spec = SPEC_EWALD_HALO + OpAdapter for partial S(k)
energy:	E_Ewald (replicated)
forces:	F_Ewald (per-rank owned)

>]; out [label=<

energy = E_MACE + E_Ewald (replicated) forces = F_MACE + F_Ewald (per-rank owned)

>]; input -> padded [label=<build halo (once)>]; padded -> mace; padded -> ewald; mace -> out; ewald -> out; } ``` ```{code-block} python :caption: Pipeline construction is one line per block plus a top-level wrap. from nvalchemi.models import PipelineModelWrapper from nvalchemi.distributed import DistributedPipelineModel pipeline = PipelineModelWrapper([ MACEWrapper.from_checkpoint("medium-0b2"), EwaldModelWrapper(cutoff=10.0), ]) dist_model = DistributedPipelineModel(pipeline, domain_config) energy_dict = dist_model(sharded_batch) # energy_dict["energy"] ← MACE + Ewald summed, globally replicated # energy_dict["forces"] ← per-rank owned, autograd-derived ``` Composition rules at the seam: | Sub-model A | Sub-model B | Pipeline strategy | |---|---|---| | Halo | Halo | Halo (single padded_batch, both blocks read it) | | Halo | Sharded | Sharded (most permissive) | | Sharded | Sharded | Sharded | | Local | anything | the other one | The merge rule is implemented in the module-level `_merge_policies` helper in `nvalchemi/distributed/spec.py` (with `_merge_compile_policies` for the compile contract), driven by `MLIPSpec.__or__` — same discriminated-union pattern as the Strategy classes themselves. --- ## 7. Wrapping Warp / Triton kernels ### 7.1 The boundary problem `ShardTensor.__torch_function__` only fires on ops PyTorch dispatches through the public Python API. Warp / Triton kernels reach into tensor data via `wp.from_torch(t)` / Triton's pointer protocol — both strip the subclass before reading. The kernel sees a plain buffer and writes a plain buffer; ShardTensor never gets a chance to intervene. ```{graphviz} :caption: An OpAdapter wraps the kernel boundary. Inputs enter as ShardTensors; the adapter pre-shapes them per the wrapper's declared transforms (e.g. slice to owned only); the kernel runs on plain tensors; outputs get post-shaped (e.g. shell-rows-back-to-owners) and re-promoted to ShardTensor for the rest of the model. :align: center digraph kernel_boundary { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; inp [label=<

input: ShardTensor (n_padded, 3)
owned	shell

>]; pre [label=<

pre-shape: e.g. slice to owned-only (arg_transforms = SliceOwned)

owned (plain tensor)

>]; kern [label=<

kernel runs Warp / Triton / custom_op plain tensors only — ShardTensor is invisible inside

>]; post [label=<

post-shape: e.g. send shell partials to owners, refresh shell from owners (output_transforms = ScatterOutputs)
owned (corrected)	shell (refreshed)

>]; out [label=<

output: ShardTensor again, ready for the next op
owned	shell

>]; inp -> pre [label=<unwrap + transform>]; pre -> kern [label=<launch>]; kern -> post [label=<cross-rank correction>]; post -> out [label=<re-promote to ShardTensor>]; } ``` ### 7.2 The transform vocabulary Every input / output transform is a small dataclass marker. The framework's `wrap_custom_op` interprets them at call time. | Position | Transform | What it does | |---|---|---| | input | `GatherInputs` | halo-pad an owned-shape input to `(n_padded, *F)` | | input | `GatherInputsFull` | sharded analogue: full-gather to `(n_global+1, *F)` | | input | `SliceOwned` | slice halo-padded input to `(n_owned, *F)` | | output | `ScatterOutputs` | halo_reverse + halo_forward on a per-atom output | | output | `AllReduceSum` | cross-rank SUM (autograd-symmetric) | | output | `SliceOutputsOwned` | slice global-shape output back to owned-only | ### 7.3 Worked example ```{code-block} python :caption: A Warp pair-energy kernel wrapped through ``OpAdapter``. (Excerpted from ``examples/distributed/05_byo_graph_transformer.py``.) @wp.kernel def _gaussian_pair_kernel(...): ... @torch.library.custom_op("tutorial::gaussian_pair_energy", mutates_args=()) def gaussian_pair_energy(edge_index, positions, epsilon, sigma, cutoff): energy_per_atom = torch.zeros(...) wp.launch(_gaussian_pair_kernel, ...) return energy_per_atom # Spec declares the boundary semantics. spec = MLIPSpec( distribution=DistributionSpec( policy=HaloStoragePolicy(), custom_ops=( OpAdapter( op=torch.ops.tutorial.gaussian_pair_energy.default, arg_transforms={}, # halo-padded inputs OK as-is output_transforms={0: ScatterOutputs()}, # output[0] is per-atom: halo-correct ), ), ), output_kinds={"energy": OutputKind.PER_GRAPH, ...}, ) ``` The OpAdapter is the *only* distribution-aware code in the wrapper. The kernel itself stays single-process; the spec parameterises the cross-rank wrap. --- ## 8. Validation + Bring-Your-Own-Model ### 8.1 The flow `trace_and_validate` is the BYO author's only required entry point. A single call: build a sample, point at the model factory, get back a verdict + a working spec. ```{graphviz} :caption: ``trace_and_validate`` flow. A single reference run captures the truth; ``world_size`` workers re-run the same factory with a candidate spec; diffs that exceed tolerance trigger the auto-fix engine, which proposes a spec mutation and retries. :align: center digraph validate { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; Factory [label=<

model_factory() a callable that returns a fresh wrapper — called once for the reference, once per spawned worker

>]; Ref [label=<

reference run single-process forward on the sample batch — produces the per-output truth tensors plus an op-trace and helper-trace

>]; Spec0 [label=<

initial candidate spec use ``wrapper.distribution_spec`` if declared, otherwise infer a sensible halo default

>]; Spawn [label=<

spawn ``world_size`` workers each runs the wrapper through the framework with the candidate spec

>]; Diff [label=<

compare to reference per-output abs/rel diff, op firings, halo completeness, helper-diagnostic gaps

>]; Pass [label=<

report.ok = True save spec, ship alongside checkpoint

>]; AutoFix [label=<

auto-fix rule engine try a known mutation: • swap halo correction → local • promote per-graph autograd output to all-reduce • drop a redundant all-reduce

>]; Translate [label=<

error translator rewrite generic torch errors as framework-specific hints (e.g. dropped scatter return)

>]; Factory -> Ref -> Spec0; Spec0 -> Spawn -> Diff; Diff -> Pass [label=<diff < tolerance>]; Diff -> AutoFix [label=<diff > tolerance>]; Diff -> Translate [label=<worker raised>]; AutoFix -> Spawn [label=<retry with mutated spec>]; } ``` ### 8.2 What the report carries ```{code-block} python :caption: The actionable surface. Either ``report.ok`` is True and ``report.spec`` is ready to save, or ``report.next_action`` tells you exactly what's wrong. report = trace_and_validate(model_factory, sample_batch, world_size=2) if report.ok: report.spec.save("my_model_spec.json") else: report.log_summary(logger) # Output includes: # - validation status + auto-fix applied # - per-output abs/rel diffs vs single-process # - dispatch-handler firings (so you can see what the multi-rank # run actually exercised) # - halo-completeness verdict # - helper-diagnostic gaps from watched third-party packages # - "Diagnosis:" hint when an error pattern is recognised # (e.g. dropped scatter_add_ return, missing OpAdapter, etc.) ``` ### 8.3 The intended user path ```{graphviz} :caption: The BYO arc — the same five steps regardless of whether the model is pure PyTorch (example 04) or has a Warp kernel (example 05). Most users finish at step 5 without ever touching step 4. :align: center digraph byo { rankdir=TB; node [shape=plaintext fontname="Helvetica"]; S1 [label=<

1. write the single-process wrapper no distribution code, no spec — just BaseModelMixin

>]; S2 [label=<

2. call trace_and_validate no arguments beyond ``model_factory`` + sample batch — auto-fix discovers the spec for typical halo MPNNs and per-rank-partial outputs

>]; S3 [label=<

3. read report.log_summary(logger) if it passed, the report shows the residual diff vs single-process; if it failed, the diagnostic points at the root cause in plain English

>]; S4 [label=<

4. (rare) author an OpAdapter only needed if the model embeds a Warp / Triton kernel: declare the input pre-shape and output post-shape on the spec

>]; S5 [label=<

5. spec.save("model_spec.json") ship alongside the checkpoint; production loads it as ``MLIPSpec.load(...)``

>]; S1 -> S2 -> S3; S3 -> S4 [label=<opaque kernel?>]; S3 -> S5 [label=<otherwise>]; S4 -> S2 [label=<re-validate>]; } ``` ```{code-block} python :caption: The end-to-end happy path is six lines. (Excerpted from ``examples/distributed/04_byo_pytorch_mpnn.py``.) def model_factory(): torch.manual_seed(123) return BPWrapper(BPModel(feat_dim=32, cutoff=5.0)).cuda() report = trace_and_validate(model_factory, sample_batch, world_size=2) report.log_summary(logger) # validation PASSED in 1 attempt report.spec.save("bp_model_spec.json") # Production: # spec = MLIPSpec.load("bp_model_spec.json") # dist = DistributedModel(BPWrapper(BPModel()), domain_cfg, spec=spec) ``` --- ## What's not in this overview * **Performance numbers** — see `examples/distributed/benchmark_*.py` and the scaling tables those produce. The benchmarks measure per-step wall clock, halo-build amortisation, and weak/strong scaling on argon / sodium chloride / silica supercells. * **Checkpoint compatibility** — covered in {doc}`distributed_byo`. * **The full handler registry** — every predicate + handler is documented in {doc}`distributed_shardtensor` (this overview only walks the dispatch flow at the conceptual level). * **Failure-mode catalogue** — common spec mistakes and the diagnostics that catch them are in {doc}`distributed_byo` § "Common failure modes". For a runnable end-to-end build, work through the two BYO examples in order: | Example | Adds | |---|---| | `examples/distributed/04_byo_pytorch_mpnn.py` | The minimal pure-PyTorch path. | | `examples/distributed/05_byo_graph_transformer.py` | The Warp-kernel path with a hand-authored `OpAdapter`. |