Route isend/irecv to nonblocking backend methods and stage them as async on MPS

deepspeed/comm/comm.py

@@ -384,13 +384,13 @@ def recv(tensor, src=None, group=None, tag=0, prof=False, log_name='recv', debug
 @timed_op
 def isend(tensor, dst, group=None, tag=0, prof=False, log_name='isend', debug=get_caller_func()):
     global cdb
-    return cdb.send(tensor=tensor, dst=dst, group=group, tag=tag)
+    return cdb.isend(tensor=tensor, dst=dst, group=group, tag=tag)
 
 
 @timed_op
 def irecv(tensor, src=None, group=None, tag=0, prof=False, log_name='irecv', debug=get_caller_func()):
     global cdb
-    return cdb.recv(tensor=tensor, src=src, group=group, tag=tag)
+    return cdb.irecv(tensor=tensor, src=src, group=group, tag=tag)
 
 
 @timed_op

deepspeed/comm/torch.py

@@ -103,23 +103,29 @@ def __init__(self, work, copy_back):
         self.copy_back = copy_back
 
     def wait(self):
+        result = None
         if self.work is not None:
-            self.work.wait()
+            result = self.work.wait()
         self.copy_back()
-        return None
+        return result
 
 
 def _needs_cpu_staging(tensor):
     # gloo (the only torch backend on macOS) cannot operate on MPS tensors.
     return isinstance(tensor, torch.Tensor) and tensor.device.type == 'mps'
 
 
-def stage_on_cpu(func):
+def stage_on_cpu(func=None, *, always_async=False):
     """Runs a collective on CPU copies of any MPS tensor arguments, then copies the results back.
 
     This is what lets DeepSpeed use the gloo backend on Apple Silicon, where device tensors are
     not supported by any torch.distributed backend. Unified memory keeps the copies cheap.
+
+    always_async is for ops like isend/irecv that are asynchronous by contract but have no
+    async_op parameter: the copy back must wait until the returned work handle completes.
     """
+    if func is None:
+        return lambda wrapped_func: stage_on_cpu(wrapped_func, always_async=always_async)
 
     def _stage(arg, pairs):
         if _needs_cpu_staging(arg):
@@ -147,7 +153,7 @@ def copy_back():
 
         # async_op is usually forwarded positionally, so resolve it against the real signature.
         bound_args = signature.bind(self, *args, **kwargs)
-        if bound_args.arguments.get('async_op', False):
+        if always_async or bound_args.arguments.get('async_op', False):
             return StagedWork(work, copy_back)
         copy_back()
         return work
@@ -427,12 +433,12 @@ def recv(self, tensor, src=None, group=None, tag=0):
         return torch.distributed.recv(tensor=tensor, src=src, group=group, tag=tag)
 
     @disable_compiler_collective
-    @stage_on_cpu
+    @stage_on_cpu(always_async=True)
     def isend(self, tensor, dst, group=None, tag=0):
         return torch.distributed.isend(tensor=tensor, dst=dst, group=group, tag=tag)
 
     @disable_compiler_collective
-    @stage_on_cpu
+    @stage_on_cpu(always_async=True)
     def irecv(self, tensor, src=None, group=None, tag=0):
         return torch.distributed.irecv(tensor=tensor, src=src, group=group, tag=tag)
 

op_builder/cpu/comm.py

@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import os
+import sys
 from .builder import CPUOpBuilder
 
 
@@ -65,6 +66,7 @@ def cxx_args(self):
         return ['-O2', '-fopenmp']
 
     def is_compatible(self, verbose=False):
-        # TODO: add soft compatibility check for private binary release.
-        #  a soft check, as in we know it can be trivially changed.
+        # The shared-memory kernels use Linux-only APIs, so let other platforms fall back to gloo.
+        if sys.platform != 'linux':
+            return False
         return super().is_compatible(verbose)

tests/unit/comm/test_dist.py

@@ -128,6 +128,94 @@ def test(self, num_elements):
         assert torch.all(x == result)
 
 
+class FakeP2PWork:
+    """Stands in for a torch.distributed Work so the staging wrapper can be tested single-process."""
+
+    def __init__(self, fill=None):
+        self.fill = fill
+
+    def wait(self):
+        if self.fill is not None:
+            self.fill()
+        return True
+
+
+@pytest.mark.skipif(get_accelerator().device_name() != 'mps', reason="covers the MPS CPU-staging path")
+class TestMpsStagedP2P:
+
+    def test_irecv_defers_copy_back_to_wait(self, monkeypatch):
+        from deepspeed.comm.torch import TorchBackend, StagedWork
+        captured = {}
+
+        def fake_irecv(tensor, src=None, group=None, tag=0):
+            captured['staged'] = tensor
+            return FakeP2PWork(fill=lambda: tensor.copy_(torch.arange(16, dtype=torch.float32)))
+
+        monkeypatch.setattr(torch.distributed, 'irecv', fake_irecv)
+        backend = TorchBackend.__new__(TorchBackend)
+        received = torch.zeros(16, dtype=torch.float32, device='mps')
+
+        handle = backend.irecv(received, src=0)
+
+        # gloo must see a CPU tensor, and the MPS tensor must stay untouched until wait().
+        assert isinstance(handle, StagedWork)
+        assert captured['staged'].device.type == 'cpu'
+        assert received.abs().sum().item() == 0
+        assert handle.wait() is True
+        assert torch.equal(received.cpu(), torch.arange(16, dtype=torch.float32))
+
+    def test_isend_stages_payload_on_cpu(self, monkeypatch):
+        from deepspeed.comm.torch import TorchBackend, StagedWork
+        captured = {}
+
+        def fake_isend(tensor, dst=None, group=None, tag=0):
+            captured['staged'] = tensor
+            return FakeP2PWork()
+
+        monkeypatch.setattr(torch.distributed, 'isend', fake_isend)
+        backend = TorchBackend.__new__(TorchBackend)
+        payload = torch.arange(16, dtype=torch.float32, device='mps')
+
+        handle = backend.isend(payload, dst=0)
+
+        assert isinstance(handle, StagedWork)
+        assert captured['staged'].device.type == 'cpu'
+        assert torch.equal(captured['staged'], torch.arange(16, dtype=torch.float32))
+        assert handle.wait() is True
+
+
+class TestDistIsendIrecv(DistributedTest):
+    world_size = 2
+
+    def _launch_procs(self, num_procs, init_method):
+        # Two gloo ranks do not need two devices, but the base class gates process count on
+        # device_count(), which reports sockets on CPU and would skip this test on CI. Bypass
+        # the gate there and pin gloo so oneCCL bindings are not silently picked up.
+        if get_accelerator().device_name() != 'cpu':
+            return super()._launch_procs(num_procs, init_method)
+        self.backend = 'gloo'
+        torch.multiprocessing.set_start_method('forkserver', force=True)
+        self._launch_daemonic_procs(num_procs, init_method)
+
+    def test(self):
+        rank = dist.get_rank()
+        device = get_accelerator().device_name()
+        length = 4096
+        if rank == 0:
+            payload = torch.arange(length, dtype=torch.float32).to(device)
+            handle = dist.isend(payload, dst=1)
+        else:
+            received = torch.zeros(length, dtype=torch.float32).to(device)
+            handle = dist.irecv(received, src=0)
+        # isend/irecv are asynchronous by contract: they must hand back a waitable handle,
+        # and the received buffer is only valid after wait() completes.
+        assert hasattr(handle, 'wait')
+        handle.wait()
+        if rank == 1:
+            expected = torch.arange(length, dtype=torch.float32).to(device)
+            assert torch.equal(received, expected)
+
+
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("num_elements", [128, 3])
 class TestDistInferenceAllReduce(DistributedTest):