Route isend/irecv to nonblocking backend methods and stage them as async on MPS (#8303)

## Summary

Resolves @delock's review note on #8293
(#8293 (comment)):
`irecv` is asynchronous by contract and has no `async_op` parameter, so
the MPS CPU-staging wrapper must handle it explicitly.

Two fixes:

1. **`deepspeed/comm/comm.py`** — `isend`/`irecv` dispatched to the
*blocking* `cdb.send`/`cdb.recv` (since the original comm backend,
#1985). Callers got a blocking call and `recv`'s return value (the
source rank `int`) instead of a waitable handle, so
`dist.irecv(...).wait()` raised `AttributeError`. This affects every
backend, not just MPS — e.g. the 1-bit comm helpers
(`runtime/comm/{compressed,hccl,nccl}.py`) call
`dist.isend/irecv(...).wait()`. They now route to
`cdb.isend`/`cdb.irecv`.
2. **`deepspeed/comm/torch.py`** — with the routing fixed, the MPS
staging wrapper's copy-back decision (keyed on an `async_op` argument)
ran immediately for `irecv`, before the transfer completed. A new
`always_async` flag on `stage_on_cpu` defers the copy-back to the
handle's `wait()` for `isend`/`irecv`. `StagedWork.wait()` now also
returns the underlying work's wait result.

### Verified (M5 Max, macOS 26.3, torch 2.13)

- Real two-process gloo run with MPS tensors: on master, `dist.irecv`
returns an `int` and `.wait()` crashes; with this PR it returns a handle
and the buffer holds the correct payload after `wait()`.
- `DS_ACCELERATOR=mps pytest unit/comm/test_dist.py`: 10 passed
(multi-rank cases skip on 1 device).
- ZeRO-2/3 smoke training unaffected.

### Test

Adds `TestDistIsendIrecv` (world size 2) to the existing
`tests/unit/comm/test_dist.py`: rank 0 `isend`s, rank 1 `irecv`s, both
assert a waitable handle and verify the payload after `wait()`.
Backend-agnostic, so it exercises the routing fix on CUDA/CPU CI as
well.

### Relation to #8301

#8301 addresses the same note with a more extensive `StagedWork`
(futures, result identity restoration, weakref buffer tracking). This PR
makes the fix more concise and accurate: no current DeepSpeed users
calls `Work.result()`/`get_future()` on staged P2P ops, and the staged
CPU buffer for `isend` is kept alive by the deferred copy-back closure
until `wait()`. Huge Credit to @FU-max-boop for the thorough analysis of
the Work semantics and fixes.

---------

Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>

Author: PKUWZP

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):