Restore gp test

src/fairchem/core/calculate/pretrained_mlip.py

@@ -74,6 +74,7 @@ def get_predict_unit(
     cache_dir: str = CACHE_DIR,
     workers: int = 1,
     seed: int = 41,
+    gp_config=None,
 ) -> MLIPPredictUnit:
     """
     Retrieves a prediction unit for a specified model.
@@ -116,6 +117,7 @@ def get_predict_unit(
         form_elem_refs,
         workers,
         seed,
+        gp_config=gp_config,
     )
 
 

src/fairchem/core/common/gp_utils.py

@@ -8,13 +8,17 @@
 from __future__ import annotations
 
 import contextlib
+import dataclasses
 import logging
 import threading
+from dataclasses import dataclass
 
 import torch
 from torch import distributed as dist
 from torch.distributed.nn.functional import all_reduce, reduce_scatter
 
+from fairchem.core.common.utils import StrEnum
+
 """
 Functions to support graph parallel training.
 This is based on the Megatron-LM implementation:
@@ -26,6 +30,26 @@
 _GRAPH_PARALLEL_GROUP = None
 _DATA_PARALLEL_GROUP = None
 
+
+class GPMode(StrEnum):
+    ALLGATHER = "allgather"
+    ALL_TO_ALL = "all_to_all"
+
+
+class GPPartition(StrEnum):
+    INDEX_SPLIT = "index_split"
+    SPATIAL = "spatial"
+
+
+@dataclass
+class GraphParallelConfig:
+    group_size: int = 1
+    mode: GPMode = GPMode.ALLGATHER
+    partition: GPPartition = GPPartition.INDEX_SPLIT
+
+
+_GP_CONFIG: GraphParallelConfig | None = None
+
 _tls = threading.local()
 
 
@@ -97,6 +121,13 @@ def setup_graph_parallel_groups(
         if i == found[0]:
             _GRAPH_PARALLEL_GROUP = group
 
+    # Ensure a GP config exists so downstream code can read
+    # `get_gp_config().mode` without a None check. Callers that want
+    # non-default settings (A2A, spatial partition) should call
+    # `set_gp_config` explicitly after this.
+    if _GP_CONFIG is None:
+        set_gp_config(GraphParallelConfig(group_size=graph_parallel_group_size))
+
 
 def setup_gp(config) -> None:
     gp_size = config["gp_gpus"]
@@ -129,10 +160,18 @@ def setup_gp(config) -> None:
         if i == found[0]:
             _GRAPH_PARALLEL_GROUP = group
 
+    # Every entry point that sets up GP groups must also set a GP config so
+    # downstream code (e.g. escn_md.py) can read `get_gp_config().mode`
+    # without a None check. setup_graph_parallel_groups()'s callers set the
+    # config alongside; do the same here for parity.
+    if _GP_CONFIG is None:
+        set_gp_config(GraphParallelConfig(group_size=gp_size))
+
 
 def cleanup_gp() -> None:
     global _DATA_PARALLEL_GROUP
     global _GRAPH_PARALLEL_GROUP
+    global _GP_CONFIG
     assert _GRAPH_PARALLEL_GROUP is not None
     assert _DATA_PARALLEL_GROUP is not None
     with contextlib.suppress(ValueError):
@@ -141,12 +180,57 @@ def cleanup_gp() -> None:
         dist.destroy_process_group(_GRAPH_PARALLEL_GROUP)
     _DATA_PARALLEL_GROUP = None
     _GRAPH_PARALLEL_GROUP = None
+    _GP_CONFIG = None
 
 
 def initialized() -> bool:
     return _GRAPH_PARALLEL_GROUP is not None
 
 
+def set_gp_config(config: GraphParallelConfig) -> None:
+    global _GP_CONFIG
+    _GP_CONFIG = config
+
+
+def get_gp_config() -> GraphParallelConfig | None:
+    return _GP_CONFIG
+
+
+def resolve_gp_config_for_workers(
+    gp_config: GraphParallelConfig | None,
+    num_workers: int,
+) -> GraphParallelConfig | None:
+    """
+    Reconcile a user-provided GraphParallelConfig with the target number
+    of workers.
+
+    Behavior:
+      - ``num_workers <= 1``: no GP is used; return ``gp_config`` unchanged
+        (may be ``None``).
+      - ``num_workers > 1``:
+          * If ``gp_config`` is ``None``, build a default
+            ``GraphParallelConfig(group_size=num_workers)``.
+          * If the config's ``group_size`` is still the default (1),
+            return a copy with ``group_size=num_workers``. The caller's
+            config is NOT mutated.
+          * If ``group_size == num_workers`` already, return it unchanged.
+          * Otherwise raise ``ValueError`` — an explicit mismatch is a
+            configuration error.
+    """
+    if num_workers <= 1:
+        return gp_config
+    if gp_config is None:
+        return GraphParallelConfig(group_size=num_workers)
+    if gp_config.group_size == 1:
+        return dataclasses.replace(gp_config, group_size=num_workers)
+    if gp_config.group_size != num_workers:
+        raise ValueError(
+            f"gp_config.group_size ({gp_config.group_size}) must equal "
+            f"num_workers ({num_workers})"
+        )
+    return gp_config
+
+
 def get_dp_group():
     return _DATA_PARALLEL_GROUP
 

src/fairchem/core/common/parallelism/__init__.py

@@ -12,6 +12,7 @@
     GPContext,
     all_to_all_collect,
     build_gp_context,
+    compute_a2a_partition,
 )
 from fairchem.core.common.parallelism.graph_partition import (
     PartitionStrategy,
@@ -25,6 +26,7 @@
     "PartitionStrategy",
     "all_to_all_collect",
     "build_gp_context",
+    "compute_a2a_partition",
     "partition_atoms_index_split",
     "partition_atoms_spatial",
 ]

src/fairchem/core/common/parallelism/graph_parallel_a2a.py

@@ -14,6 +14,11 @@
 from torch.profiler import record_function
 
 from fairchem.core.common import gp_utils
+from fairchem.core.common.parallelism.graph_partition import (
+    PartitionStrategy,
+    partition_atoms_index_split,
+    partition_atoms_spatial,
+)
 
 
 def _safe_all_to_all(
@@ -61,8 +66,6 @@ class GPContext:
     Graph parallel context holding communication metadata for all-to-all.
 
     Runtime-only struct: every field is needed for the forward/backward pass.
-    Construction intermediates (node_partition, rank_assignments, needed_atoms,
-    global_to_local, etc.) are computed in build_gp_context but not stored.
 
     Attributes:
         rank: Current GP rank.
@@ -79,6 +82,7 @@ class GPContext:
         total_recv: Total number of embeddings to receive (sum of recv_splits).
         local_edge_idx: Indices into edge_index_local where source is a local atom.
         remote_edge_idx: Indices into edge_index_local where source is a remote atom.
+        rank_assignments: Rank owner for each atom, shape (total_atoms,).
     """
 
     rank: int
@@ -93,6 +97,7 @@ class GPContext:
     total_recv: int
     local_edge_idx: torch.Tensor
     remote_edge_idx: torch.Tensor
+    rank_assignments: torch.Tensor
 
 
 def _sparse_index_exchange(
@@ -191,6 +196,41 @@ def _sparse_index_exchange(
     return send_counts, send_indices_global
 
 
+@torch.compiler.disable
+def compute_a2a_partition(
+    pos: torch.Tensor,
+    total_atoms: int,
+    device: torch.device,
+    world_size: int,
+    rank: int,
+    strategy: PartitionStrategy,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute rank assignments and local node partition for A2A graph parallel.
+
+    Args:
+        pos: Atom positions, shape (N, 3).
+        total_atoms: Total number of atoms.
+        device: Device for output tensors.
+        world_size: Number of GP ranks.
+        rank: Current GP rank.
+        strategy: Partitioning strategy (SPATIAL or INDEX_SPLIT).
+
+    Returns:
+        Tuple of (rank_assignments, node_partition) where rank_assignments
+        is shape (N,) mapping each atom to a rank, and node_partition is
+        the indices of atoms belonging to this rank.
+    """
+    with record_function("a2a_partition"):
+        if strategy == PartitionStrategy.SPATIAL:
+            rank_assignments = partition_atoms_spatial(pos, world_size)
+        else:
+            rank_assignments = partition_atoms_index_split(
+                total_atoms, world_size, device
+            )
+    node_partition = (rank_assignments == rank).nonzero(as_tuple=True)[0]
+    return rank_assignments, node_partition
+
+
 @torch.compiler.disable
 def build_gp_context(
     edge_index: torch.Tensor,
@@ -313,6 +353,7 @@ def build_gp_context(
         total_recv=total_recv,
         local_edge_idx=local_edge_idx,
         remote_edge_idx=remote_edge_idx,
+        rank_assignments=rank_assignments,
     )
 
 
@@ -524,7 +565,6 @@ def backward(ctx, grad_received: torch.Tensor):
 def all_to_all_collect(
     x_local: torch.Tensor,
     gp_ctx: GPContext,
-    send_indices: torch.Tensor,
 ) -> torch.Tensor:
     """
     High-level function to collect remote embeddings via all-to-all.
@@ -536,15 +576,14 @@ def all_to_all_collect(
     Args:
         x_local: Local atom embeddings, shape (local_atoms, *features).
         gp_ctx: Graph parallel context.
-        send_indices: Local indices of atoms to send.
 
     Returns:
         x_received: Remote atom embeddings,
             shape (total_needed, *features).
     """
     return AllToAllCollect.apply(
         x_local,
-        send_indices,
+        gp_ctx.send_indices,
         gp_ctx.send_counts,
         gp_ctx.recv_counts,
         gp_utils.get_gp_group(),