feat: add distributed training execution

docs/perlmutter.md

@@ -0,0 +1,58 @@
+# Perlmutter execution
+
+The package does not encode Perlmutter paths, modules, or allocation policy.
+Create the uv environment in the project location appropriate for your
+account, select a site-compatible GPU/driver environment, and use the same
+Hydra configuration as on a workstation.
+
+## Single GPU
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=/path/to/graphs.pt \
+  environment=perlmutter \
+  trainer.device=cuda \
+  trainer.max_epochs=10
+```
+
+The equivalent Slurm wrapper is:
+
+```bash
+sbatch scripts/slurm/train_single_gpu.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+## Single-node DDP
+
+```bash
+GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_gpu.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+The wrapper uses `torchrun`; `batch_size` and `num_workers` are per GPU.
+Effective batch size is `data.batch_size * number_of_processes`, and only rank
+0 writes the shared checkpoint/config/prediction artifacts.
+
+## Multi-node DDP
+
+```bash
+GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_node.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+Use the provided script as a template and adapt only allocation/account
+settings required by the site. Evaluation and prediction can use
+`scripts/slurm/evaluate.sh`; prediction currently gathers moderate-size
+results in memory.
+
+## Checks and common failures
+
+```bash
+uv run python -c "import torch, dgl; print(torch.__version__, dgl.__version__, torch.cuda.is_available())"
+uv run gnn4colliders --help
+```
+
+An unavailable DGL wheel, incompatible driver, missing cache, or mismatched
+cache schema should be fixed in the environment/input rather than hidden with
+package-level path changes. GPU kernels and distributed execution are not
+promised bitwise deterministic.

src/gnn4colliders/config/application.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from pathlib import Path
 from typing import Any
 
@@ -10,6 +11,7 @@ import torch
 from omegaconf import DictConfig, OmegaConf
 
 from gnn4colliders.data import (
+    DistributedGraphDataLoader,
     GraphDataLoader,
     GraphDataset,
     GraphSample,
@@ -18,6 +20,12 @@ from gnn4colliders.data import (
     SplitDefinition,
     select_split,
 )
+from gnn4colliders.distributed import (
+    DistributedContext,
+    barrier,
+    finalize,
+    initialize,
+)
 from gnn4colliders.features import build_node_features
 from gnn4colliders.graphs import build_dgl_graph
 from gnn4colliders.inference import Predictor, write_npz
@@ -55,7 +63,9 @@ def load_samples(config: DictConfig) -> GraphDataset:
     return GraphDataset(cache.load())
 
 
-def loaders(config: DictConfig) -> dict[str, GraphDataLoader]:
+def loaders(
+    config: DictConfig, context: DistributedContext | None = None
+) -> dict[str, GraphDataLoader]:
     samples = load_samples(config)
     split = SplitDefinition(
         train_folds=frozenset(config.data.splits.train_folds),
@@ -65,22 +75,41 @@ def loaders(config: DictConfig) -> dict[str, GraphDataLoader]:
     result = {}
     for name in ("train", "validation", "test"):
         selected = GraphDataset(select_split(samples.samples, split, name))
-        result[name] = GraphDataLoader(
-            selected,
-            int(config.data.batch_size),
-            shuffle=bool(config.data.shuffle) if name == "train" else False,
-            seed=int(config.data.seed),
+        loader_type = (
+            DistributedGraphDataLoader
+            if context is not None and context.enabled
+            else GraphDataLoader
+        )
+        result[name] = (
+            loader_type(
+                selected,
+                int(config.data.batch_size),
+                context,
+                shuffle=bool(config.data.shuffle) if name == "train" else False,
+                seed=int(config.data.seed),
+            )
+            if loader_type is DistributedGraphDataLoader
+            else loader_type(
+                selected,
+                int(config.data.batch_size),
+                shuffle=bool(config.data.shuffle) if name == "train" else False,
+                seed=int(config.data.seed),
+            )
         )
     return result
 
 
 def prepare(config: DictConfig) -> Path:
+    if bool(config.distributed.enabled) or int(os.environ.get("WORLD_SIZE", "1")) > 1:
+        raise ValueError(
+            "prepare does not support distributed launch; run it once on rank 0"
+        )
     data = config.data
     if not data.files:
         raise ValueError("prepare requires data.files")
-    feature_branches = data.get("feature_branches")
-    object_types = data.get("object_types")
-    scales = data.get("scales")
+    feature_branches = plain(data.get("feature_branches"))
+    object_types = plain(data.get("object_types"))
+    scales = plain(data.get("scales"))
     if feature_branches is None or object_types is None or scales is None:
         raise ValueError(
             "prepare requires data.feature_branches, object_types, and scales"
@@ -90,8 +119,7 @@ def prepare(config: DictConfig) -> Path:
         tree_name=str(data.tree_name),
         label=data.get("label", 1),
         feature_branches=feature_branches,
-        tracking_info=data.get("tracking_info", []),
-        global_features=data.get("global_features", []),
+        global_features=plain(data.get("global_features", [])),
         fold_var=str(data.get("fold_var", "eventNumber")),
         weight_var=data.get("weight_var"),
     )
@@ -131,57 +159,124 @@ def _model_and_task(config: DictConfig, loader: GraphDataLoader):
 
 def train(config: DictConfig) -> Path:
     validate_config(plain(config))
-    seed_everything(int(config.trainer.seed))
-    run_root = Path(config.environment.output_root)
-    save_resolved(config, run_root)
-    split_loaders = loaders(config)
-    model, task = _model_and_task(config, split_loaders["train"])
-    trainer = build_trainer(plain(config.trainer), model=model, task=task)
-    if config.checkpoint.resume:
-        payload = CheckpointManager.load(config.checkpoint.resume, map_location="cpu")
-        restore_training_state(
-            payload,
-            model=model,
-            trainer=trainer,
+    context = initialize(
+        enabled=True if config.distributed.enabled else None,
+        backend=config.distributed.backend,
+        device=config.trainer.device,
+    )
+    try:
+        seed_everything(int(config.trainer.seed) + context.rank)
+        run_root = Path(config.environment.output_root)
+        if context.is_main_process:
+            save_resolved(config, run_root)
+        barrier(context)
+        split_loaders = loaders(config, context)
+        model, task = _model_and_task(config, split_loaders["train"])
+        trainer = build_trainer(
+            plain(config.trainer), model=model, task=task, distributed_context=context
+        )
+        if config.checkpoint.resume:
+            payload = CheckpointManager.load(
+                config.checkpoint.resume, map_location="cpu"
+            )
+            restore_training_state(
+                payload,
+                model=trainer.model,
+                trainer=trainer,
+                optimizer=trainer.optimizer,
+                scheduler=trainer.scheduler,
+                early_stopping=trainer.early_stopping,
+            )
+        validation_loader = (
+            split_loaders["validation"] if len(split_loaders["validation"]) else None
+        )
+        trainer.fit(
+            split_loaders["train"],
+            validation_loader,
+            epochs=int(config.trainer.max_epochs),
+        )
+        if not context.is_main_process:
+            barrier(context)
+            return (
+                Path(config.checkpoint.directory)
+                / f"epoch_{trainer.state.epoch:04d}.pt"
+            )
+        manager = CheckpointManager(config.checkpoint.directory)
+        path = manager.save(
+            model=trainer.model,
+            trainer_state=trainer.state,
             optimizer=trainer.optimizer,
             scheduler=trainer.scheduler,
             early_stopping=trainer.early_stopping,
+            task_config=plain(config.task),
+            model_config=plain(config.model),
         )
-    validation_loader = (
-        split_loaders["validation"] if len(split_loaders["validation"]) else None
-    )
-    trainer.fit(
-        split_loaders["train"], validation_loader, epochs=int(config.trainer.max_epochs)
-    )
-    manager = CheckpointManager(config.checkpoint.directory)
-    return manager.save(
-        model=model,
-        trainer_state=trainer.state,
-        optimizer=trainer.optimizer,
-        scheduler=trainer.scheduler,
-        early_stopping=trainer.early_stopping,
-        task_config=plain(config.task),
-        model_config=plain(config.model),
-    )
+        barrier(context)
+        return path
+    finally:
+        finalize(context)
 
 
 def predict_or_evaluate(config: DictConfig, *, evaluate: bool = False) -> Any:
     validate_config(plain(config))
+    context = initialize(
+        enabled=True if config.distributed.enabled else None,
+        backend=config.distributed.backend,
+        device=config.trainer.device,
+    )
+    try:
+        split_loaders = loaders(config, context)
+        model, task = _model_and_task(config, split_loaders[config.inference.split])
+        payload = CheckpointManager.load(
+            config.inference.checkpoint, map_location="cpu"
+        )
+        from gnn4colliders.training import load_model_weights
+
+        load_model_weights(model, payload)
+        predictor = Predictor(
+            model,
+            task,
+            device=str(config.trainer.device),
+            distributed_context=context,
+        )
+        result = (
+            predictor.evaluate(split_loaders[config.inference.split])
+            if evaluate
+            else predictor.predict(split_loaders[config.inference.split])
+        )
+        if not evaluate and context.is_main_process:
+            output = Path(config.inference.output)
+            if str(config.inference.format).lower() != "npz":
+                raise ValueError("only npz inference output is currently supported")
+            write_npz(result, output)
+        barrier(context)
+        return result
+    finally:
+        finalize(context)
+
+
+def export_onnx(config: DictConfig) -> Path:
+    """Export a checkpoint using the first prepared graph batch."""
+    if bool(config.distributed.enabled) or int(os.environ.get("WORLD_SIZE", "1")) > 1:
+        raise ValueError("ONNX export must run on a single process")
+    checkpoint = config.export.checkpoint
+    if not checkpoint:
+        raise ValueError("export.checkpoint is required")
+    if str(config.export.format).lower() != "onnx":
+        raise ValueError("only export.format=onnx is supported")
     split_loaders = loaders(config)
-    model, task = _model_and_task(config, split_loaders[config.inference.split])
-    payload = CheckpointManager.load(config.inference.checkpoint, map_location="cpu")
-    from gnn4colliders.training import load_model_weights
-
-    load_model_weights(model, payload)
-    predictor = Predictor(model, task, device=str(config.trainer.device))
-    result = (
-        predictor.evaluate(split_loaders[config.inference.split])
-        if evaluate
-        else predictor.predict(split_loaders[config.inference.split])
+    split = str(config.export.split)
+    if split not in split_loaders:
+        raise ValueError(f"unsupported export split {split!r}")
+    example_batch = next(iter(split_loaders[split]), None)
+    if example_batch is None:
+        raise ValueError(f"export split {split!r} is empty")
+    from gnn4colliders.export import export_checkpoint_to_onnx
+
+    return export_checkpoint_to_onnx(
+        checkpoint,
+        config.export.output,
+        example_batch=example_batch,
+        opset=int(config.export.opset),
+        overwrite=bool(config.export.overwrite),
     )
-    if not evaluate:
-        output = Path(config.inference.output)
-        if str(config.inference.format).lower() != "npz":
-            raise ValueError("only npz inference output is currently supported")
-        write_npz(result, output)
-    return result

src/gnn4colliders/distributed/__init__.py

@@ -0,0 +1,19 @@
+"""Minimal PyTorch DDP support for GNN4Colliders applications."""
+
+from .collectives import barrier, broadcast_bool, gather_objects, gather_tensor
+from .context import DistributedContext, finalize, initialize, process_group
+from .model import prepare_model
+from .sampler import DistributedIndices
+
+__all__ = [
+    "DistributedContext",
+    "DistributedIndices",
+    "barrier",
+    "broadcast_bool",
+    "finalize",
+    "gather_objects",
+    "gather_tensor",
+    "initialize",
+    "prepare_model",
+    "process_group",
+]

src/gnn4colliders/distributed/collectives.py

@@ -0,0 +1,52 @@
+"""Small tensor/object collectives used by training and inference."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+from .context import DistributedContext
+
+
+def broadcast_bool(value: bool, context: DistributedContext, *, src: int = 0) -> bool:
+    if not context.enabled:
+        return value
+    tensor = torch.tensor([int(value)], device=context.device)
+    dist.broadcast(tensor, src=src)
+    return bool(tensor.item())
+
+
+def gather_tensor(tensor: torch.Tensor, context: DistributedContext) -> torch.Tensor:
+    """All-gather a first-dimension-sharded tensor, including unequal shards."""
+
+    if not context.enabled:
+        return tensor
+    local = tensor.contiguous()
+    size = torch.tensor([local.shape[0]], device=local.device, dtype=torch.long)
+    sizes = [torch.zeros_like(size) for _ in range(context.world_size)]
+    dist.all_gather(sizes, size)
+    max_size = max(int(item.item()) for item in sizes)
+    padded = torch.zeros(
+        (max_size, *local.shape[1:]), dtype=local.dtype, device=local.device
+    )
+    padded[: local.shape[0]] = local
+    gathered = [torch.empty_like(padded) for _ in range(context.world_size)]
+    dist.all_gather(gathered, padded)
+    return torch.cat(
+        [part[: int(length.item())] for part, length in zip(gathered, sizes)], dim=0
+    )
+
+
+def gather_objects(value: Any, context: DistributedContext) -> list[Any]:
+    if not context.enabled:
+        return [value]
+    gathered: list[Any] = [None] * context.world_size
+    dist.all_gather_object(gathered, value)
+    return gathered
+
+
+def barrier(context: DistributedContext) -> None:
+    if context.enabled:
+        dist.barrier()

src/gnn4colliders/distributed/context.py

@@ -0,0 +1,103 @@
+"""Process-group context and rank-local device selection."""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Iterator
+
+import torch
+import torch.distributed as dist
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    enabled: bool = False
+    rank: int = 0
+    local_rank: int = 0
+    world_size: int = 1
+    device: torch.device = torch.device("cpu")
+
+    @property
+    def is_main_process(self) -> bool:
+        return self.rank == 0
+
+    @classmethod
+    def single_process(cls, device: torch.device | str = "cpu") -> "DistributedContext":
+        return cls(device=torch.device(device))
+
+
+def _launched() -> bool:
+    return "RANK" in os.environ or "WORLD_SIZE" in os.environ
+
+
+def initialize(
+    *,
+    enabled: bool | None = None,
+    backend: str | None = None,
+    device: torch.device | str | None = None,
+    timeout_seconds: int = 1800,
+) -> DistributedContext:
+    """Initialize ``torchrun``'s ``env://`` process group when requested."""
+
+    launched = _launched()
+    if enabled is False and launched and int(os.environ.get("WORLD_SIZE", "1")) > 1:
+        raise ValueError(
+            "distributed launcher variables are present but distributed.enabled=false"
+        )
+    active = launched if enabled is None else enabled
+    if not active:
+        requested = torch.device(
+            device or ("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        if requested.type == "cuda" and not torch.cuda.is_available():
+            requested = torch.device("cpu")
+        return DistributedContext.single_process(requested)
+
+    try:
+        rank = int(os.environ["RANK"])
+        local_rank = int(os.environ.get("LOCAL_RANK", rank))
+        world_size = int(os.environ["WORLD_SIZE"])
+    except KeyError as error:
+        raise ValueError(
+            "distributed execution requires torchrun RANK and WORLD_SIZE"
+        ) from error
+    if world_size < 2:
+        return DistributedContext.single_process(device or "cpu")
+    requested = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    if requested.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError(
+                "CUDA was requested for distributed execution but is unavailable"
+            )
+        torch.cuda.set_device(local_rank)
+        assigned = torch.device("cuda", local_rank)
+        selected_backend = backend or "nccl"
+    else:
+        assigned = torch.device("cpu")
+        selected_backend = backend or "gloo"
+    if not dist.is_initialized():
+        dist.init_process_group(
+            backend=selected_backend,
+            init_method="env://",
+            timeout=timedelta(seconds=timeout_seconds),
+        )
+    return DistributedContext(True, rank, local_rank, world_size, assigned)
+
+
+def finalize(context: DistributedContext) -> None:
+    """Destroy a process group created for ``context``."""
+
+    if context.enabled and dist.is_available() and dist.is_initialized():
+        dist.destroy_process_group()
+
+
+@contextmanager
+def process_group(**kwargs: object) -> Iterator[DistributedContext]:
+    context = initialize(**kwargs)
+    try:
+        yield context
+    finally:
+        finalize(context)

src/gnn4colliders/distributed/model.py

@@ -0,0 +1,20 @@
+"""DDP model wrapping kept outside architecture implementations."""
+
+from __future__ import annotations
+
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+
+from .context import DistributedContext
+
+
+def prepare_model(model: nn.Module, context: DistributedContext) -> nn.Module:
+    model = model.to(context.device)
+    if not context.enabled:
+        return model
+    kwargs = (
+        {"device_ids": [context.local_rank], "output_device": context.local_rank}
+        if context.device.type == "cuda"
+        else {}
+    )
+    return DistributedDataParallel(model, **kwargs)

src/gnn4colliders/distributed/sampler.py

@@ -0,0 +1,47 @@
+"""Index sharding for the project's deterministic graph loader."""
+
+from __future__ import annotations
+
+import torch
+
+from .context import DistributedContext
+
+
+class DistributedIndices:
+    """Epoch-aware rank partitioning without evaluation padding duplicates."""
+
+    def __init__(
+        self,
+        size: int,
+        context: DistributedContext,
+        *,
+        shuffle: bool,
+        drop_last: bool = False,
+        seed: int = 0,
+        pad: bool = True,
+    ) -> None:
+        self.size = size
+        self.context = context
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.seed = seed
+        self.pad = pad
+        self.epoch = 0
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+    def indices(self) -> list[int]:
+        values = torch.arange(self.size, dtype=torch.long)
+        if self.shuffle:
+            generator = torch.Generator().manual_seed(self.seed + self.epoch)
+            values = values[torch.randperm(self.size, generator=generator)]
+        if not self.context.enabled:
+            return values.tolist()
+        if self.drop_last:
+            usable = self.size - self.size % self.context.world_size
+            values = values[:usable]
+        elif self.pad and self.size % self.context.world_size:
+            extra = self.context.world_size - self.size % self.context.world_size
+            values = torch.cat((values, values[:extra]))
+        return values[self.context.rank :: self.context.world_size].tolist()

tests/unit/distributed/test_ddp_cpu.py

@@ -0,0 +1,76 @@
+import multiprocessing as mp
+import os
+import socket
+
+import pytest
+import torch
+from torch import nn
+
+from gnn4colliders.distributed import finalize, initialize, prepare_model
+
+
+def _port() -> int:
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _worker(rank: int, port: int, queue) -> None:
+    os.environ.update(
+        RANK=str(rank),
+        LOCAL_RANK=str(rank),
+        WORLD_SIZE="2",
+        MASTER_ADDR="127.0.0.1",
+        MASTER_PORT=str(port),
+    )
+    context = initialize(enabled=True, backend="gloo", device="cpu")
+    try:
+        torch.manual_seed(5)
+        model = prepare_model(nn.Linear(1, 1), context)
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+        x = torch.tensor([[1.0], [2.0]]) if rank == 0 else torch.tensor([[3.0], [4.0]])
+        y = 2 * x
+        optimizer.zero_grad()
+        loss = nn.functional.mse_loss(model(x), y)
+        loss.backward()
+        optimizer.step()
+        queue.put(
+            (
+                rank,
+                model.module.weight.detach().cpu().item(),
+                model.module.bias.detach().cpu().item(),
+            )
+        )
+    finally:
+        finalize(context)
+
+
+@pytest.mark.skipif(
+    not torch.distributed.is_available(), reason="torch.distributed unavailable"
+)
+def test_two_rank_cpu_ddp_matches_single_process_update():
+    queue = mp.get_context("spawn").Queue()
+    # Both workers must share one rendezvous port.
+    port = _port()
+    processes = [
+        mp.get_context("spawn").Process(target=_worker, args=(rank, port, queue))
+        for rank in range(2)
+    ]
+    for process in processes:
+        process.start()
+    values = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+        assert process.exitcode == 0
+    assert values[0][1] == values[1][1]
+    assert values[0][2] == values[1][2]
+
+    torch.manual_seed(5)
+    reference = nn.Linear(1, 1)
+    optimizer = torch.optim.SGD(reference.parameters(), lr=0.1)
+    x = torch.arange(1.0, 5.0).reshape(-1, 1)
+    optimizer.zero_grad()
+    nn.functional.mse_loss(reference(x), 2 * x).backward()
+    optimizer.step()
+    assert torch.allclose(torch.tensor(values[0][1]), reference.weight.squeeze())
+    assert torch.allclose(torch.tensor(values[0][2]), reference.bias.squeeze())