fix: keep DD communication inside domain subgroups

nvalchemi/distributed/_core/gather_primitives.py

@@ -49,16 +49,16 @@
 
 
 def mesh_group(mesh: Any) -> Any:
-    """Return the default ``ProcessGroup`` for *mesh*.
+    """Return the concrete ``ProcessGroup`` for *mesh*.
 
-    Accepts either a real ``DeviceMesh`` or a test-harness mock; returns
-    ``None`` for both "no distribution configured" and "mesh present but not
-    group-capable".
+    Accepts a concrete ``ProcessGroup`` unchanged, or resolves one from a real
+    ``DeviceMesh`` or test-harness mock. Returns ``None`` for both "no
+    distribution configured" and "mesh present but not group-capable".
 
     Parameters
     ----------
-    mesh : DeviceMesh or object or None
-        The device mesh to resolve a group from.
+    mesh : ProcessGroup, DeviceMesh, object, or None
+        The process group to preserve or device mesh to resolve.
 
     Returns
     -------
@@ -68,6 +68,8 @@ def mesh_group(mesh: Any) -> Any:
     """
     if mesh is None:
         return None
+    if isinstance(mesh, dist.ProcessGroup):
+        return mesh
     get_group = getattr(mesh, "get_group", None)
     if get_group is None:
         return None
@@ -100,15 +102,18 @@ def _funcol_group_arg(mesh: Any) -> Any:
     Prefers the ``(DeviceMesh, 0)`` spec — the form Dynamo special-cases (and
     the one physicsnemo's own compiled collectives use). Real distributed
     inference always carries a real ``DeviceMesh``, so the compiled path takes
-    this branch. Falls back to the resolved ``ProcessGroup`` only for eager
-    test harnesses that pass a lightweight (non-``DeviceMesh``) mesh; the
-    ``isinstance`` is a compile-time constant, so under ``torch.compile`` only
-    the traceable branch survives.
+    this branch. A concrete ``ProcessGroup`` is returned unchanged. Otherwise,
+    falls back to the resolved ``ProcessGroup`` for eager test harnesses that
+    pass a lightweight (non-``DeviceMesh``) mesh; the ``isinstance`` checks are
+    compile-time constants, so under ``torch.compile`` only the traceable branch
+    survives.
     """
     from torch.distributed.device_mesh import DeviceMesh  # noqa: PLC0415
 
     if isinstance(mesh, DeviceMesh):
         return (mesh, 0)
+    if isinstance(mesh, dist.ProcessGroup):
+        return mesh
     return funcol_group(mesh)
 
 
@@ -343,18 +348,24 @@ def _neighbor_p2p_v_1d(
             if send_counts[r] > 0:
                 recv[r_off[r] : r_off[r + 1]].copy_(send[s_off[r] : s_off[r + 1]])
             continue
+        peer = dist.get_global_rank(group, r) if group is not None else r
         if send_counts[r] > 0:
             ops.append(
                 dist.P2POp(
                     dist.isend,
                     send[s_off[r] : s_off[r + 1]].contiguous(),
-                    r,
+                    peer,
                     group=group,
                 )
             )
         if recv_counts[r] > 0:
             ops.append(
-                dist.P2POp(dist.irecv, recv[r_off[r] : r_off[r + 1]], r, group=group)
+                dist.P2POp(
+                    dist.irecv,
+                    recv[r_off[r] : r_off[r + 1]],
+                    peer,
+                    group=group,
+                )
             )
     if ops:
         for work in dist.batch_isend_irecv(ops):
@@ -421,8 +432,11 @@ def _neighbor_p2p_fixed(
         if r == rank:
             continue
         sl = slice(r * m, (r + 1) * m)
-        ops.append(dist.P2POp(dist.isend, send_rows[sl].contiguous(), r, group=group))
-        ops.append(dist.P2POp(dist.irecv, recv[sl], r, group=group))
+        peer = dist.get_global_rank(group, r) if group is not None else r
+        ops.append(
+            dist.P2POp(dist.isend, send_rows[sl].contiguous(), peer, group=group)
+        )
+        ops.append(dist.P2POp(dist.irecv, recv[sl], peer, group=group))
     if ops:
         for work in dist.batch_isend_irecv(ops):
             work.wait()
@@ -467,7 +481,7 @@ def halo_exchange_fixed(
             is_nccl = False
         if is_nccl:
             return _neighbor_p2p_fixed(send_rows, world_size, neighbors, group)
-    return funcol_all_to_all_fixed(send_rows, world_size, None)
+    return funcol_all_to_all_fixed(send_rows, world_size, group)
 
 
 def funcol_all_to_all_v_rows(
@@ -581,10 +595,11 @@ def _isend_irecv_v_1d(
         if r == rank:
             recv_slice.copy_(send_slice)
         else:
+            peer = dist.get_global_rank(group, r) if group is not None else r
             if send_slice.numel() > 0:
-                ops.append(dist.isend(send_slice, dst=r, group=group))
+                ops.append(dist.isend(send_slice, dst=peer, group=group))
             if recv_slice.numel() > 0:
-                ops.append(dist.irecv(recv_slice, src=r, group=group))
+                ops.append(dist.irecv(recv_slice, src=peer, group=group))
     for op in ops:
         op.wait()
 

nvalchemi/distributed/_core/particle_halo.py

@@ -55,6 +55,19 @@
 
 logger = logging.getLogger(__name__)
 
+# ``torch.library.custom_op`` schemas cannot carry a Python ``ProcessGroup``.
+# Publish the active halo group immediately before each model forward so the
+# eager custom-op bodies can resolve their collectives to the same domain group
+# as the surrounding DD strategy. ``None`` preserves the default-group behavior
+# for direct/single-mesh callers.
+_HALO_PROCESS_GROUP: Any = None
+
+
+def set_halo_process_group(group: Any) -> None:
+    """Publish the domain process group used by compiled halo custom ops."""
+    global _HALO_PROCESS_GROUP
+    _HALO_PROCESS_GROUP = group
+
 
 # ======================================================================
 # Ghost identification
@@ -725,19 +738,21 @@ def _halo_a2a_v_default_group(
     rank: int,
     world_size: int,
 ) -> torch.Tensor:
-    """:func:`_funcol_indexed_all_to_all_v_rows` over the DEFAULT process group.
+    """:func:`_funcol_indexed_all_to_all_v_rows` over the active halo group.
 
     Used inside the halo-correction custom op, which runs eagerly at runtime and
-    cannot take a ``DeviceMesh`` arg. Valid for single-domain-mesh-dim topology
-    (the domain group IS the world group); ``funcol`` mesh=None resolves the
-    default group.
+    cannot take a ``DeviceMesh`` or ``ProcessGroup`` arg. The surrounding halo
+    strategy publishes its domain group before the model forward; direct callers
+    that publish no group retain the default-world behavior.
     """
     send_rows = torch.cat(
         [tensor.index_select(0, indices[j]) for j in range(world_size)], dim=0
     )
     send_counts = [int(sizes[rank][j]) for j in range(world_size)]
     recv_counts = [int(sizes[i][rank]) for i in range(world_size)]
-    return funcol_all_to_all_v_rows(send_rows, send_counts, recv_counts, None)
+    return funcol_all_to_all_v_rows(
+        send_rows, send_counts, recv_counts, _HALO_PROCESS_GROUP
+    )
 
 
 def _halo_scatter_correct_dense(
@@ -749,7 +764,7 @@ def _halo_scatter_correct_dense(
     world_size: int,
 ) -> torch.Tensor:
     """``halo_forward_exchange(halo_reverse_exchange(padded))`` as pure
-    compute + collective (no autograd.Function, default group)."""
+    compute + collective (no autograd.Function, active halo group)."""
     # reverse: fold borrowed halo rows back into their owners.
     halo = padded[n_owned:].contiguous()
     rev_indices: list[torch.Tensor] = []
@@ -892,10 +907,10 @@ def halo_forward_op(
     this rank's halo (ghost) region.
 
     Compile-safe counterpart of :func:`halo_forward_exchange`; runs eagerly at
-    runtime (default group), while the trace sees only the registered fake. Its
-    adjoint (backward) is :func:`halo_scatter_correct_op`. The marker arrays ride
-    as ``int[]`` (not tensors) so inductor lowering does not see a real-Tensor
-    constant alongside the fake ``owned`` input.
+    runtime on the active halo group, while the trace sees only the registered
+    fake. Its adjoint (backward) is :func:`halo_scatter_correct_op`. The marker
+    arrays ride as ``int[]`` (not tensors) so inductor lowering does not see a
+    real-Tensor constant alongside the fake ``owned`` input.
 
     Parameters
     ----------
@@ -1199,10 +1214,10 @@ def halo_forward_static_op(
     world_size: int,
 ) -> torch.Tensor:
     """Fixed-shape owned->[owned|ghost] refresh. Runs eagerly at runtime
-    (default group); the trace sees only the fake. Owned rows pass through;
+    on the active halo group; the trace sees only the fake. Owned rows pass through;
     ghost rows are gathered from neighbors via a uniform-split all_to_all."""
     send_rows = padded_in.index_select(0, send_index)
-    recv = halo_exchange_fixed(send_rows, world_size, None)
+    recv = halo_exchange_fixed(send_rows, world_size, _HALO_PROCESS_GROUP)
     recv = recv * _row_mask_like(recv_real, recv).to(recv.dtype)
     ghost_acc = torch.zeros_like(padded_in).index_add(0, recv_dest, recv)
     rowidx = torch.arange(padded_in.shape[0], device=padded_in.device)
@@ -1232,7 +1247,7 @@ def _hfs_backward(ctx, grad_out):  # type: ignore[no-untyped-def]
     grad_recv = grad_ghost.index_select(0, recv_dest) * _row_mask_like(
         recv_real, grad_out
     ).to(grad_out.dtype)
-    grad_send = halo_exchange_fixed(grad_recv, ws, None)
+    grad_send = halo_exchange_fixed(grad_recv, ws, _HALO_PROCESS_GROUP)
     grad_in = grad_out * (~ghostmask).to(grad_out.dtype)
     grad_in = grad_in.index_add(0, send_index, grad_send)
     return grad_in, None, None, None, None, None
@@ -1262,13 +1277,13 @@ def halo_scatter_correct_static_op(
 
     # reverse: ghost rows (recv-slot order) -> owning rank -> index_add into owners.
     ghost_rows = padded_in.index_select(0, recv_dest) * recv_real_f
-    back = halo_exchange_fixed(ghost_rows, world_size, None)
+    back = halo_exchange_fixed(ghost_rows, world_size, _HALO_PROCESS_GROUP)
     owned_only = (padded_in * (~ghostmask).to(padded_in.dtype)).to(acc_dt)
     owned_acc = owned_only.index_add(0, send_index, back.to(acc_dt)).to(padded_in.dtype)
 
     # forward: re-broadcast corrected owners to ghosts.
     send_rows = owned_acc.index_select(0, send_index)
-    recv = halo_exchange_fixed(send_rows, world_size, None) * recv_real_f
+    recv = halo_exchange_fixed(send_rows, world_size, _HALO_PROCESS_GROUP) * recv_real_f
     ghost_acc = torch.zeros_like(padded_in).index_add(0, recv_dest, recv)
     return torch.where(ghostmask, ghost_acc, owned_acc)
 

nvalchemi/distributed/_dynamics_coordinator.py

@@ -256,8 +256,8 @@ def _restore(self) -> None:
 
     def broadcast_state(self, batch: Batch) -> None:
         """Broadcast the replicated controller + cell state (the integrator's
-        declared ``__dd_replicated__``) from rank 0 so floating-point divergence
-        cannot accumulate over a long run.
+        declared ``__dd_replicated__``) from group-local rank 0 so floating-point
+        divergence cannot accumulate over a long run.
 
         With global KE/DOF and identical config the controller evolves
         identically on every rank, so this is anti-drift insurance, not a
@@ -266,7 +266,7 @@ def broadcast_state(self, batch: Batch) -> None:
         if not self.active:
             return
         group = self._strategy.process_group
-        src = 0
+        src = dist.get_global_rank(group, 0) if group is not None else 0
         state = getattr(self._dyn, "_state", None)
         fields = getattr(self._dyn, "__dd_replicated__", ())
         for name in fields:

nvalchemi/distributed/sharded_batch.py

@@ -270,7 +270,7 @@ def rank_assignment(self) -> torch.Tensor:
         # Single all_gather into a flat (world_size,) tensor, one sync.
         n_owned_t = torch.tensor([self.n_owned], dtype=torch.int64, device=device)
         sizes_t = torch.empty(world_size, dtype=torch.int64, device=device)
-        dist.all_gather_into_tensor(sizes_t, n_owned_t)
+        dist.all_gather_into_tensor(sizes_t, n_owned_t, group=mesh_group(self.mesh))
 
         # Build the block-constant assignment via repeat_interleave — no
         # per-rank Python loop or slicing.

nvalchemi/distributed/strategy.py

@@ -52,6 +52,7 @@
     mesh_group,
     set_halo_neighbor_ranks,
 )
+from nvalchemi.distributed._core.particle_halo import set_halo_process_group
 
 if TYPE_CHECKING:
     from nvalchemi.data.batch import Batch
@@ -307,6 +308,9 @@ def run_forward(
         """Run the model forward on this rank's halo-padded (owned + ghost) shard."""
         dist_model._dist_ctx.cap_atoms = self.caps_atoms
         dist_model._dist_ctx.strategy = self
+        # Compiled halo custom ops cannot carry a Python ProcessGroup through
+        # their dispatcher schema, so publish this strategy's exact domain group.
+        set_halo_process_group(self.process_group)
         return _halo_run_forward(dist_model, state, wired_fields)
 
     def on_cell_change(self, state: ShardState, cell: torch.Tensor | None) -> None: