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.

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.#

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:

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 ℓ.#

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.

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.#

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.

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.#

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.

Dispatch flowchart for a single torch op. Predicates inspect the inputs’ shapes + spec; the matching handler runs cross-rank logic and re-promotes outputs.#

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:

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 Bring Your Own Model: Authoring a Distribution Spec 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 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.

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.#

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.#

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.

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.#

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.

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.#

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

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#

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.#

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

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.#

4.3 Why this is short to write#

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

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:

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:

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.

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.#

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.

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.#

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#

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.

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.#

8.2 What the report carries#

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#

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.#

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 Bring Your Own Model: Authoring a Distribution Spec.
The full handler registry — every predicate + handler is documented in ShardTensor: How Per-Atom Fields Flow Across Ranks (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 Bring Your Own Model: Authoring a Distribution Spec § “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`.