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)
 

tests/unit/comm/test_dist.py

@@ -128,6 +128,28 @@ def test(self, num_elements):
         assert torch.all(x == result)
 
 
+class TestDistIsendIrecv(DistributedTest):
+    world_size = 2
+
+    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):