refit clean up and refactoring

megatron/core/resharding/copy_services/base.py

@@ -2,9 +2,29 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Optional
 
 import torch
+import torch.distributed as dist
+
+
+@dataclass
+class SendOp:
+    """Single send operation pending in a CopyService queue."""
+
+    task_id: int | None
+    tensor: torch.Tensor
+    dest_rank: int
+
+
+@dataclass
+class RecvOp:
+    """Single receive operation pending in a CopyService queue."""
+
+    task_id: int | None
+    tensor: torch.Tensor
+    src_rank: int
 
 
 class CopyService(ABC):
@@ -17,6 +37,13 @@ class CopyService(ABC):
     remote transfers simply ignore it.
     """
 
+    def __init__(self, group=None):
+        self.group = group
+        # group.rank()/size() supports cross-cluster ProcessGroups where members
+        # have independent default PGs.
+        self.rank = group.rank() if group is not None else dist.get_rank()
+        self.world_size = group.size() if group is not None else dist.get_world_size()
+
     @abstractmethod
     def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optional[int] = None):
         """Register a tensor send from the current rank to ``dest_rank``."""
@@ -31,3 +58,45 @@ def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optiona
     def run(self):
         """Execute all previously submitted send/recv operations as a single batch."""
         ...
+
+    def close(self) -> None:
+        """Release backend-owned resources.  Default no-op; NVSHMEM overrides."""
+
+
+def match_local_ops_by_task_id(
+    local_sends: list, local_recvs: list, backend_name: str, rank: int
+) -> list[tuple]:
+    """Pair same-rank send/recv ops by task_id, raising on any mismatch.
+
+    Returns a list of ``(send_op, recv_op)`` tuples for the caller to apply
+    backend-specific local-copy logic.  Either op type may be a backend-local
+    wrapper as long as it exposes ``.task_id``.
+    """
+    sends_by_id = {op.task_id: op for op in local_sends}
+    recvs_by_id = {op.task_id: op for op in local_recvs}
+    if None in sends_by_id or None in recvs_by_id:
+        raise RuntimeError(
+            f"{backend_name}: local (same-rank) transfer requires a task_id "
+            "to match sends with recvs"
+        )
+    # Count mismatch catches both imbalanced send/recv lists (which would
+    # otherwise silently drop the longer side) and duplicate task_ids (which
+    # collapse to fewer dict entries than list entries).
+    if (
+        len(local_sends) != len(local_recvs)
+        or len(sends_by_id) != len(local_sends)
+        or len(recvs_by_id) != len(local_recvs)
+    ):
+        raise RuntimeError(
+            f"{backend_name}: unmatched local ops on rank {rank}: "
+            f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs"
+        )
+    pairs = []
+    for task_id, recv_op in recvs_by_id.items():
+        send_op = sends_by_id.get(task_id)
+        if send_op is None:
+            raise RuntimeError(
+                f"{backend_name}: missing local send for task_id={task_id} on rank {rank}"
+            )
+        pairs.append((send_op, recv_op))
+    return pairs

megatron/core/resharding/copy_services/gloo_copy_service.py

@@ -2,52 +2,35 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
 from typing import List, Optional, Tuple
 
 import torch
 import torch.distributed as dist
 
-from .base import CopyService
+from .base import CopyService, RecvOp, SendOp, match_local_ops_by_task_id
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class SendOp:
-    """Simple container describing a single send operation."""
-
-    task_id: int | None
-    tensor: torch.Tensor
-    dest_rank: int
-
-
-@dataclass
-class RecvOp:
-    """Simple container describing a single receive operation."""
-
-    task_id: int | None
-    tensor: torch.Tensor
-    src_rank: int
-
-
 class GlooCopyService(CopyService):
     """
     CopyService implementation that routes refit traffic over a CPU/Gloo
     process group instead of NCCL.
     """
 
     def __init__(self, group=None):
+        super().__init__(group=group)
         if group is not None:
             self.gloo_pg = group
-            self.rank = group.rank()
-            self.world_size = group.size()
         else:
-            self.rank = dist.get_rank()
-            self.world_size = dist.get_world_size()
             self.gloo_pg = dist.new_group(backend="gloo")
         self.send_ops: List[SendOp] = []
+        # Each recv op is paired with its GPU destination tensor; the SendOp/RecvOp
+        # itself carries a pinned-CPU staging buffer for Gloo's CPU PG.
         self.recv_ops: List[Tuple[RecvOp, torch.Tensor]] = []
+        # Dedicated stream for GPU-side work (local same-rank copies and the
+        # final CPU->GPU writebacks) so they overlap with the GPU->CPU staging
+        # copies issued on the default stream during ``run()``.
         self._copy_stream = torch.cuda.Stream()
         if self.rank == 0:
             logger.info(
@@ -74,43 +57,23 @@ def run(self):
                 f"{len(self.send_ops)} sends + {len(self.recv_ops)} recvs = {total_ops} ops"
             )
 
-        p2p_ops: List[dist.P2POp] = []
-
-        # Short-circuit self transfers into local device copies.
         local_sends = [op for op in self.send_ops if op.dest_rank == self.rank]
         remote_sends = [op for op in self.send_ops if op.dest_rank != self.rank]
         local_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank == self.rank]
         remote_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank != self.rank]
 
+        # Local copies run on a dedicated stream so they overlap with the
+        # GPU->CPU staging copies issued on the default stream below.
         if local_sends or local_recvs:
-            local_sends_by_id = {op.task_id: op for op in local_sends}
-            if None in local_sends_by_id:
-                raise RuntimeError(
-                    "GlooCopyService: local (same-rank) transfer requires a task_id "
-                    "to match sends with recvs"
-                )
-            local_recvs_by_id = {recv.task_id: (recv, dst) for (recv, dst) in local_recvs}
-            if None in local_recvs_by_id:
-                raise RuntimeError(
-                    "GlooCopyService: local (same-rank) transfer requires a task_id "
-                    "to match sends with recvs"
-                )
-            if len(local_sends_by_id) != len(local_sends) or len(local_recvs_by_id) != len(
-                local_recvs
-            ):
-                raise RuntimeError(
-                    f"GlooCopyService: unmatched local ops on rank {self.rank}: "
-                    f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs"
-                )
-            for task_id, (recv_op, dst_tensor) in local_recvs_by_id.items():
-                send_op = local_sends_by_id.get(task_id)
-                if send_op is None:
-                    raise RuntimeError(
-                        f"GlooCopyService: missing local send for task_id={task_id} "
-                        f"on rank {self.rank}"
-                    )
-                with torch.no_grad():
+            local_recv_objs = [recv for recv, _ in local_recvs]
+            dst_by_task_id = {recv.task_id: dst for recv, dst in local_recvs}
+            pairs = match_local_ops_by_task_id(
+                local_sends, local_recv_objs, "GlooCopyService", self.rank
+            )
+            with torch.no_grad(), torch.cuda.stream(self._copy_stream):
+                for send_op, recv_op in pairs:
                     src_tensor = send_op.tensor
+                    dst_tensor = dst_by_task_id[recv_op.task_id]
                     if dst_tensor.device != src_tensor.device:
                         dst_tensor.copy_(src_tensor.to(dst_tensor.device))
                     else:
@@ -128,13 +91,15 @@ def run(self):
             )
             cpu_tensor.copy_(op.tensor.detach(), non_blocking=True)
             cpu_send_bufs.append(cpu_tensor)
-        # Single sync after all GPU→CPU copies are issued.
+        # Wait only on default-stream staging copies; _copy_stream keeps running.
         if cpu_send_bufs:
-            torch.cuda.synchronize()
+            torch.cuda.current_stream().synchronize()
 
         for op in remote_sends:
+            # Drop the GPU reference now that staging is complete.
             op.tensor = None
 
+        p2p_ops: List[dist.P2POp] = []
         for cpu_tensor, op in zip(cpu_send_bufs, remote_sends):
             p2p_ops.append(
                 dist.P2POp(dist.isend, cpu_tensor, group=self.gloo_pg, group_peer=op.dest_rank)
@@ -150,18 +115,17 @@ def run(self):
                 req.wait()
 
         # Copy received CPU buffers back into the original destination tensors.
-        # Use non_blocking with pinned memory for overlap.
-        for recv, dst_tensor in remote_recvs:
-            if dst_tensor.is_cuda:
-                dst_tensor.copy_(recv.tensor, non_blocking=True)
-            else:
-                dst_tensor.copy_(recv.tensor)
-
-        if self._copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self._copy_stream)
-
-        # Ensure all async CPU→GPU copies are complete.
-        torch.cuda.synchronize()
+        # Use non_blocking with pinned memory for overlap.  Routed through
+        # _copy_stream so subsequent default-stream work isn't blocked.
+        with torch.cuda.stream(self._copy_stream):
+            for recv, dst_tensor in remote_recvs:
+                if dst_tensor.is_cuda:
+                    dst_tensor.copy_(recv.tensor, non_blocking=True)
+                else:
+                    dst_tensor.copy_(recv.tensor)
+
+        # Ensure all async CPU→GPU copies are complete and local copies have landed.
+        torch.cuda.current_stream().wait_stream(self._copy_stream)
 
         if self.rank == 0:
             logger.info("GlooCopyService: batched communication completed")

megatron/core/resharding/copy_services/nccl_copy_service.py

@@ -2,46 +2,24 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
 from typing import List, Optional
 
 import torch
 import torch.distributed as dist
 
-from .base import CopyService
+from .base import CopyService, RecvOp, SendOp, match_local_ops_by_task_id
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class SendOp:
-    """Simple container describing a single send operation."""
-
-    task_id: int | None
-    tensor: torch.Tensor
-    dest_rank: int
-
-
-@dataclass
-class RecvOp:
-    """Simple container describing a single receive operation."""
-
-    task_id: int | None
-    tensor: torch.Tensor
-    src_rank: int
-
-
 class NCCLCopyService(CopyService):
     """
     Thin wrapper around torch.distributed batch_isend_irecv to submit and execute
     a batch of point-to-point sends and recvs.
     """
 
     def __init__(self, group=None):
-        self.group = group
-        # Use group.rank()/size() to support cross-cluster ProcessGroups
-        self.rank = group.rank() if group is not None else dist.get_rank()
-        self.world_size = group.size() if group is not None else dist.get_world_size()
+        super().__init__(group=group)
         self.send_ops: List[SendOp] = []
         self.recv_ops: List[RecvOp] = []
         # Dedicated stream for local (same-rank) copies to avoid unnecessary
@@ -72,35 +50,12 @@ def run(self):
         remote_recvs = [op for op in self.recv_ops if op.src_rank != self.rank]
 
         if local_sends or local_recvs:
-            local_sends_by_id = {op.task_id: op for op in local_sends}
-            if None in local_sends_by_id:
-                raise RuntimeError(
-                    "NCCLCopyService: local (same-rank) transfer requires a task_id "
-                    "to match sends with recvs"
-                )
-            local_recvs_by_id = {op.task_id: op for op in local_recvs}
-            if None in local_recvs_by_id:
-                raise RuntimeError(
-                    "NCCLCopyService: local (same-rank) transfer requires a task_id "
-                    "to match sends with recvs"
-                )
-            if len(local_sends_by_id) != len(local_sends) or len(local_recvs_by_id) != len(
-                local_recvs
-            ):
-                raise RuntimeError(
-                    f"NCCLCopyService: unmatched local ops on rank {self.rank}: "
-                    f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs"
-                )
-            for task_id, recv_op in local_recvs_by_id.items():
-                send_op = local_sends_by_id.get(task_id)
-                if send_op is None:
-                    raise RuntimeError(
-                        f"NCCLCopyService: missing local send for task_id={task_id} "
-                        f"on rank {self.rank}"
-                    )
-                with torch.no_grad():
-                    with torch.cuda.stream(self._copy_stream):
-                        recv_op.tensor.copy_(send_op.tensor)
+            pairs = match_local_ops_by_task_id(
+                local_sends, local_recvs, "NCCLCopyService", self.rank
+            )
+            with torch.no_grad(), torch.cuda.stream(self._copy_stream):
+                for send_op, recv_op in pairs:
+                    recv_op.tensor.copy_(send_op.tensor)
 
         p2p_ops = []
         for op in remote_sends: