Bring your own PyTorch model: from architecture to a saved spec

This walkthrough takes a small interatomic potential — written in plain PyTorch — through the canonical BYO arc:

1. Architecture. A Behler-Parrinello-style energy model: a per-atom radial descriptor + an MLP energy head. Returns atomic_energies only.

2. Wrapper. A BaseModelMixin adapter — pure single-process PyTorch. No ShardTensor, no _dist_ctx, no spec.

3. Run + inspect. Build a sample, call the wrapper, look at the shapes and values it produces.

4. Validate. trace_and_validate() spawns a multi-rank run, compares per-output tensors against the single-process reference, and (if anything diverges) tells you what needs to change. The wrapper goes in unmodified — no distributed code is needed for an MPNN-halo model whose forward is just scatter-aggregations and autograd.

5. Persist. MLIPSpec.save() writes the spec the validator discovered to disk; MLIPSpec.load() reads it back. Production wrappers ship the saved JSON alongside the checkpoint.

For wrappers with non-PyTorch kernels (Warp / Triton) where the spec needs an explicit OpAdapter declaration, see Bring your own model with a Warp kernel: from architecture to spec. For the design rationale and the adapter mechanics, see docs/userguide/distributed_byo.md.

Note

Read-oriented. main() at the bottom is runnable; the sections above are illustrative.

To execute the walkthrough end-to-end (requires CUDA):

python examples/distributed/04_byo_pytorch_mpnn.py

from __future__ import annotations

from collections import OrderedDict
from pathlib import Path
from typing import Any

import torch
from torch import nn

from nvalchemi._typing import ModelOutputs
from nvalchemi.data import AtomicData, Batch
from nvalchemi.distributed.spec import MLIPSpec
from nvalchemi.models.base import (
    BaseModelMixin,
    ModelConfig,
    NeighborConfig,
    NeighborListFormat,
)

# ====================================================================
# 1. Architecture — Behler-Parrinello-style energy model
# ====================================================================
#
#   G[i] = Σ_j f(r_ij)        — per-atom radial sum over neighbours
#   E[i] = MLP(G[i], Z[i])     — per-atom energy (Z-conditioned)
#
# Returns *atomic_energies only*. Total energy and forces are computed
# at the wrapper boundary.


class BPDescriptor(nn.Module):
    """Per-atom radial descriptor: ``G[i] = Σ_j f(r_ij)``."""

    def __init__(self, feat_dim: int, cutoff: float) -> None:
        super().__init__()
        self.cutoff = cutoff
        centers = torch.linspace(0.0, cutoff, feat_dim)
        self.register_buffer("centers", centers, persistent=True)
        self.gamma = 0.5 / ((cutoff / feat_dim) ** 2)

    def forward(
        self,
        positions: torch.Tensor,  # (n_atoms, 3)
        edge_index: torch.Tensor,  # (2, n_edges) — (sender, receiver)
    ) -> torch.Tensor:
        rij = positions[edge_index[1]] - positions[edge_index[0]]
        r = torch.linalg.vector_norm(rij, dim=-1, keepdim=True)
        cutoff_envelope = (
            0.5
            * (1.0 + torch.cos(torch.pi * r / self.cutoff))
            * (r < self.cutoff).float()
        )
        edge_feats = torch.exp(-self.gamma * (r - self.centers) ** 2) * cutoff_envelope

        G = torch.zeros(
            positions.shape[0],
            edge_feats.shape[-1],
            device=positions.device,
            dtype=positions.dtype,
        )
        receivers = edge_index[1].unsqueeze(-1).expand_as(edge_feats)
        # Rebind the return: single-process ``scatter_add_`` mutates in
        # place and returns ``self`` (a no-op rebind), but under domain
        # decomposition the dispatch handler returns a *new*
        # cross-rank-corrected tensor — without the rebind ``G`` would
        # silently stay zero.
        G = G.scatter_add_(0, receivers, edge_feats)
        return G


class BPModel(nn.Module):
    """Pure-PyTorch energy model. Returns ``{"atomic_energies": (n_atoms,)}``."""

    def __init__(self, feat_dim: int = 32, cutoff: float = 5.0) -> None:
        super().__init__()
        self.descriptor = BPDescriptor(feat_dim, cutoff)
        self.element_embedding = nn.Embedding(120, 8)
        self.head = nn.Sequential(
            nn.Linear(feat_dim + 8, feat_dim),
            nn.SiLU(),
            nn.Linear(feat_dim, 1),
        )

    def forward(
        self,
        positions: torch.Tensor,
        atomic_numbers: torch.Tensor,
        edge_index: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        G = self.descriptor(positions, edge_index)
        z_emb = self.element_embedding(atomic_numbers)
        atomic_energies = self.head(torch.cat([G, z_emb], dim=-1)).squeeze(-1)
        return {"atomic_energies": atomic_energies}


# ====================================================================
# 2. Wrapper — :class:`BaseModelMixin` adapter
# ====================================================================
#
# Pure single-process: no ``ShardTensor``, no ``_dist_ctx``, no spec.
# Distribution is the framework's problem; the wrapper just describes
# the model.


class BPWrapper(nn.Module, BaseModelMixin):
    """:class:`BaseModelMixin` adapter around :class:`BPModel`.

    Boundary concerns:

    * ``model_config`` declares the active outputs + autograd inputs.
    * ``adapt_input`` extracts the COO neighbour list and drops any
      sentinel rows the producer used for padding (a single
      distribution-agnostic filter that also catches halo-receiver
      sentinels under domain decomposition).
    * ``forward`` runs the model, scatter-aggregates per-atom energies
      into per-graph totals, and computes forces via autograd.
    """

    def __init__(self, model: BPModel) -> None:
        super().__init__()
        self.model = model
        self.model_config = ModelConfig(
            outputs=frozenset({"energy", "forces"}),
            active_outputs={"energy", "forces"},
            autograd_outputs=frozenset({"forces"}),
            autograd_inputs=frozenset({"positions"}),
            neighbor_config=NeighborConfig(
                cutoff=model.descriptor.cutoff, format=NeighborListFormat.COO
            ),
        )

    @property
    def cutoff(self) -> float:
        return self.model_config.neighbor_config.cutoff

    @property
    def embedding_shapes(self) -> dict[str, tuple[int, ...]]:
        return {}

    def compute_embeddings(
        self, data: AtomicData | Batch, **kwargs: Any
    ) -> AtomicData | Batch:
        return data

    def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]:
        n_atoms = data.positions.shape[0]
        edge_index = data.neighbor_list.long().T  # (2, E)
        # Drop NL sentinel rows: both the builder's genuine padding rows
        # (atoms with fewer neighbours than ``max_neighbors``) and, under
        # domain decomposition, the halo-receiver rows the framework
        # rewrote to the same sentinel to enforce one rank per edge. One
        # line covers both regimes.
        valid = (edge_index[0] < n_atoms) & (edge_index[1] < n_atoms)
        return {
            "atomic_numbers": data.atomic_numbers,
            "edge_index": edge_index[:, valid],
        }

    def adapt_output(
        self, model_output: dict[str, torch.Tensor], data: AtomicData | Batch
    ) -> ModelOutputs:
        out: ModelOutputs = OrderedDict()
        for key in self.model_config.active_outputs:
            if key in model_output:
                out[key] = model_output[key]
        return out

    def forward(self, data: AtomicData | Batch) -> ModelOutputs:
        positions = data.positions
        compute_forces = "forces" in self.model_config.active_outputs

        if compute_forces and not positions.requires_grad:
            positions.requires_grad_(True)

        n_graphs = data.num_graphs
        kwargs = self.adapt_input(data)
        kwargs["positions"] = positions
        raw = self.model(**kwargs)
        atomic_energies = raw["atomic_energies"]

        total_energy = torch.zeros(
            n_graphs, device=positions.device, dtype=positions.dtype
        )
        total_energy = total_energy.scatter_add_(0, data.batch_idx, atomic_energies)
        raw["energy"] = total_energy

        if compute_forces:
            (forces_grad,) = torch.autograd.grad(
                outputs=total_energy.sum(),
                inputs=positions,
                create_graph=False,
            )
            raw["forces"] = -forces_grad

        return self.adapt_output(raw, data)


# ====================================================================
# 3. Run the model + inspect the outputs (single-process)
# ====================================================================


def _build_lattice_batch(n_per_side: int, device: torch.device) -> Batch:
    """Non-periodic simple-cubic cluster. Dense enough that a 5 Å NL has real
    edges. Non-periodic on purpose: this descriptor uses
    ``positions[recv] - positions[send]`` with no minimum-image shift, so it is
    only correct without periodic wraparound — under PBC the single-process
    neighbor list (shift vectors) and the domain-decomposed halo (ghost images)
    resolve cross-boundary edges differently, so the two would not agree. A cell
    is still supplied so the partitioner has a box to decompose."""
    spacing = 2.5
    coords = torch.arange(n_per_side, dtype=torch.float32, device=device) * spacing
    gx, gy, gz = torch.meshgrid(coords, coords, coords, indexing="ij")
    positions = torch.stack([gx.flatten(), gy.flatten(), gz.flatten()], dim=-1)
    n_atoms = positions.shape[0]
    torch.manual_seed(0)
    positions = positions + 0.05 * torch.randn_like(positions)
    box = n_per_side * spacing
    atomic_numbers = torch.randint(
        low=1, high=20, size=(n_atoms,), device=device, dtype=torch.long
    )
    cell = torch.eye(3, device=device, dtype=torch.float32) * box
    pbc = torch.tensor([[False, False, False]], device=device)
    data = AtomicData(
        positions=positions,
        atomic_numbers=atomic_numbers,
        cell=cell.unsqueeze(0),
        pbc=pbc,
    )
    return Batch.from_data_list([data], device=device)


def _model_factory():
    """Wrapper factory the validator uses. Called once in the
    launcher (reference) and once per spawned worker; the
    deterministic seed keeps every replica in sync.

    Set ``NVALCHEMI_BP_COMPILE=1`` to ``torch.compile`` the inner
    ``BPModel`` (the wrapper itself stays eager — only the model under
    the wrapper is compiled, mirroring ``MACEWrapper.from_checkpoint(
    compile_model=True)``). The optional ``NVALCHEMI_BP_COMPILE_BACKEND``
    overrides the backend (``inductor`` default; ``aot_eager`` for an
    Inductor-codegen-free run). Env-var-driven so the flag propagates
    cleanly through the validator's ``mp.spawn`` workers.
    """
    import os

    torch.manual_seed(123)
    model = BPModel(feat_dim=32, cutoff=5.0)
    if os.environ.get("NVALCHEMI_BP_COMPILE", "0") != "0":
        backend = os.environ.get("NVALCHEMI_BP_COMPILE_BACKEND", "inductor")
        model = torch.compile(model, backend=backend, fullgraph=False)
    return BPWrapper(model).cuda()


# ====================================================================
# 4. Validate via trace_and_validate
# 5. Persist the spec via MLIPSpec.save / load
# ====================================================================
#
# The validator returns a :class:`TraceReport` whose ``.spec`` is the
# working spec on success (or the closest variant on failure). We call
# ``.save()`` on that spec and round-trip it through ``.load()``.
# Production wrappers ship the saved JSON alongside their checkpoint.


def main() -> None:
    """Walkthrough: build → run → validate → save spec → load round-trip."""
    from loguru import logger

    if not torch.cuda.is_available():
        logger.error("This walkthrough's main() requires CUDA.")
        return

    import os

    device = torch.device("cuda:0")
    # A 2-rank partition only truly decomposes when the box is wide enough that
    # each rank's halo does NOT already cover every atom (box/ranks >
    # ~2*(cutoff+skin)); with spacing 2.5 and cutoff 5.0 that needs
    # n_per_side >= 10. The default validates on such a genuinely-decomposed
    # system (remote atoms on each rank), so trace_and_validate is real
    # evidence rather than the trivial-partition case it warns about. Override
    # with NVALCHEMI_BP_N_PER_SIDE (e.g. 4) for a quicker degenerate smoke run.
    n_per_side = int(os.environ.get("NVALCHEMI_BP_N_PER_SIDE", "10"))
    sample_batch = _build_lattice_batch(n_per_side=n_per_side, device=device)

    # ----- Section 3: run + inspect -----
    logger.info("Section 3: build, run, inspect.")
    wrapper = _model_factory()
    from nvalchemi.neighbors import compute_neighbors

    compute_neighbors(sample_batch, config=wrapper.model_config.neighbor_config)

    out = wrapper(sample_batch)
    logger.info(
        "  forward returned: "
        "energy={e:.4f} (shape {es}, dtype {ed})  |  "
        "forces shape={fs}, ‖forces‖∞={fmax:.4f}",
        e=out["energy"].item(),
        es=tuple(out["energy"].shape),
        ed=out["energy"].dtype,
        fs=tuple(out["forces"].shape),
        fmax=out["forces"].abs().max().item(),
    )

    # ----- Section 4: validate -----
    compile_state = os.environ.get("NVALCHEMI_BP_COMPILE", "0")
    backend = os.environ.get("NVALCHEMI_BP_COMPILE_BACKEND", "inductor")
    logger.info(
        "Section 4: trace_and_validate (NVALCHEMI_BP_COMPILE={c}, backend={b}).",
        c=compile_state,
        b=backend if compile_state != "0" else "n/a",
    )
    from nvalchemi.distributed.validate import trace_and_validate

    report = trace_and_validate(
        model_factory=_model_factory,
        sample_batch=sample_batch,
        world_size=2,
        device="cuda:0",
        atol=1e-4,
        rtol=1e-3,
    )
    report.log_summary(logger)

    if not report.ok:
        logger.error("Validation did not pass; not saving spec.")
        return

    # ----- Section 5: persist + round-trip -----
    spec_path = Path("bp_model_spec.json")
    report.spec.save(spec_path)
    loaded_spec = MLIPSpec.load(spec_path)
    assert loaded_spec.distribution.policy == report.spec.distribution.policy
    assert loaded_spec.output_kinds == report.spec.output_kinds
    logger.info(
        "Section 5: spec written to {p} and round-tripped cleanly. "
        "Production usage::\n"
        "    spec = MLIPSpec.load({p!r})\n"
        "    wrapper = BPWrapper(BPModel(...))\n"
        "    dist_model = DistributedModel(wrapper, domain_cfg, spec=spec)",
        p=str(spec_path),
    )


if __name__ == "__main__":
    main()