Avoid redundant copies in MPS P2P staging (#8314)

## Summary

- make staged MPS `isend` input-only: take one CPU snapshot and retain
it through completion without copying it back
- make staged MPS `irecv` output-only: allocate an empty CPU target
without reading the destination, then publish it only after the
underlying bare `wait()` succeeds
- preserve the backend's native `None` return for non-member ranks and
never publish an unconfirmed or failed receive buffer
- keep `StagedWork.wait()` deliberately narrow instead of broadening the
Work contract

Follow-up to #8303. The directional-staging scope was approved by
@delock in [this review
comment](#8303 (comment)).

## Why

#8303 made `isend` and `irecv` genuinely asynchronous and deferred MPS
publication until completion. The generic staging path still treats
every tensor as both input and output, which adds two transfers that
cannot affect a point-to-point result:

- `isend` copies the staged CPU payload back to an unchanged MPS source
after `wait()`;
- `irecv` copies the old MPS destination to CPU before the receive
overwrites it.

P2P direction is known at the call site, so those copies can be removed
without changing the default in-out staging contract used by other
collectives.

## Correctness contract

- `isend` keeps its CPU payload reachable through the returned work
handle.
- `irecv` copies CPU to MPS only after the backend `wait()` returns
successfully.
- a failed receive leaves the caller's destination unchanged.
- a backend `None` return remains `None`, including non-member group
ranks, and never publishes the empty receive target.
- `is_completed()`, timeout-aware `wait()`, `get_future()`, and
`result()` remain outside this change, matching the maintainer-approved
scope.

## Verification

Exact local head: `5caf2b6f9aef9ee52a2c3ba58a1a24059cf35add`.

- all repository pre-commit hooks for both changed files
- four single-process real-MPS direction, failure, and `None`-return
tests: `4 passed`
- real two-rank MPS/Gloo member transfer plus non-member subgroup
behavior: `2 passed`
- real two-rank CPU/Gloo member transfer plus non-member subgroup
behavior: `2 passed`
- complete `tests/unit/comm/test_dist.py` on MPS: `16 passed, 19
skipped`
- `git diff --check`

The new four-test direction/failure matrix was first run against the
unchanged #8303 source: `3 failed, 1 passed`. It is `4 passed` with this
patch.

## Local staging evidence

An instrumented no-network microbenchmark on one Apple M5 Pro host
(Python 3.12.13, PyTorch 2.13.0) confirmed the expected
Python/ATen-visible cross-device copy matrix at 1, 16, and 64 MiB:

| Operation | Creation | Wait |
|---|---|---|
| `isend` | one MPS-to-CPU snapshot | no CPU-to-MPS copy-back |
| `irecv` | no MPS-to-CPU pre-copy | one CPU-to-MPS publication after
success |

For the eliminated local phases, median p50 reductions were 99.84-99.96%
for `isend` wait and 96.41-98.11% for `irecv` creation. These
measurements isolate wrapper staging overhead with a completed fake
Work; they are not end-to-end Gloo, pipeline, training, or multi-node
speedup claims.

Signed-off-by: Fu Xiaonan <ht3fudatou@163.com>

Author: FU-max-boop

deepspeed/comm/torch.py

@@ -96,7 +96,10 @@ def wait(self):
 
 
 class StagedWork:
-    """Completes a staged collective by copying the CPU results back to the original device tensors."""
+    """Completes a staged collective after the underlying work's bare ``wait()`` succeeds.
+
+    The narrow signature is deliberate: unsupported timeout and polling APIs must fail loudly.
+    """
 
     def __init__(self, work, copy_back):
         self.work = work
@@ -115,21 +118,29 @@ def _needs_cpu_staging(tensor):
     return isinstance(tensor, torch.Tensor) and tensor.device.type == 'mps'
 
 
-def stage_on_cpu(func=None, *, always_async=False):
+def stage_on_cpu(func=None, *, always_async=False, copy_to_cpu=True, copy_from_cpu=True):
     """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.
+
+    copy_to_cpu=False allocates an output-only CPU target without reading the device tensor.
+    copy_from_cpu=False retains an input-only CPU snapshot without copying it back.
     """
     if func is None:
-        return lambda wrapped_func: stage_on_cpu(wrapped_func, always_async=always_async)
+        return lambda wrapped_func: stage_on_cpu(
+            wrapped_func,
+            always_async=always_async,
+            copy_to_cpu=copy_to_cpu,
+            copy_from_cpu=copy_from_cpu,
+        )
 
     def _stage(arg, pairs):
         if _needs_cpu_staging(arg):
-            cpu_tensor = arg.to('cpu')
+            cpu_tensor = arg.to('cpu') if copy_to_cpu else torch.empty_like(arg, device='cpu')
             pairs.append((arg, cpu_tensor))
             return cpu_tensor
         if isinstance(arg, list) and any(_needs_cpu_staging(t) for t in arg):
@@ -148,12 +159,17 @@ def wrapper(self, *args, **kwargs):
         work = func(self, *args, **kwargs)
 
         def copy_back():
-            for device_tensor, cpu_tensor in pairs:
-                device_tensor.copy_(cpu_tensor)
+            if copy_from_cpu:
+                for device_tensor, cpu_tensor in pairs:
+                    device_tensor.copy_(cpu_tensor)
 
         # async_op is usually forwarded positionally, so resolve it against the real signature.
         bound_args = signature.bind(self, *args, **kwargs)
         if always_async or bound_args.arguments.get('async_op', False):
+            # torch.distributed returns None for ranks outside the requested group. Without a
+            # completion handle, an output-only staging buffer must never be published.
+            if work is None:
+                return None
             return StagedWork(work, copy_back)
         copy_back()
         return work
@@ -433,12 +449,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(always_async=True)
+    @stage_on_cpu(always_async=True, copy_from_cpu=False)
     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(always_async=True)
+    @stage_on_cpu(always_async=True, copy_to_cpu=False)
     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

@@ -131,10 +131,13 @@ def test(self, num_elements):
 class FakeP2PWork:
     """Stands in for a torch.distributed Work so the staging wrapper can be tested single-process."""
 
-    def __init__(self, fill=None):
+    def __init__(self, fill=None, error=None):
         self.fill = fill
+        self.error = error
 
     def wait(self):
+        if self.error is not None:
+            raise self.error
         if self.fill is not None:
             self.fill()
         return True
@@ -145,55 +148,136 @@ class TestMpsStagedP2P:
 
     def test_irecv_defers_copy_back_to_wait(self, monkeypatch):
         from deepspeed.comm.torch import TorchBackend, StagedWork
-        captured = {}
+        captured = {'device_reads': 0, 'publishes': 0}
+        backend = TorchBackend.__new__(TorchBackend)
+        received = torch.zeros(16, dtype=torch.float32, device='mps')
+        real_to = torch.Tensor.to
+        real_cpu = torch.Tensor.cpu
+        real_copy = torch.Tensor.copy_
+
+        def spy_to(tensor, *args, **kwargs):
+            result = real_to(tensor, *args, **kwargs)
+            if tensor is received and result.device.type == 'cpu':
+                captured['device_reads'] += 1
+            return result
+
+        def spy_cpu(tensor, *args, **kwargs):
+            result = real_cpu(tensor, *args, **kwargs)
+            if tensor is received:
+                captured['device_reads'] += 1
+            return result
+
+        def spy_copy(tensor, source, *args, **kwargs):
+            result = real_copy(tensor, source, *args, **kwargs)
+            if source is received and tensor.device.type == 'cpu':
+                captured['device_reads'] += 1
+            if tensor is received and source.device.type == 'cpu':
+                captured['publishes'] += 1
+            return result
 
         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.Tensor, 'to', spy_to)
+        monkeypatch.setattr(torch.Tensor, 'cpu', spy_cpu)
+        monkeypatch.setattr(torch.Tensor, 'copy_', spy_copy)
         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 captured['device_reads'] == 0
+        assert captured['publishes'] == 0
         assert received.abs().sum().item() == 0
         assert handle.wait() is True
-        assert torch.equal(received.cpu(), torch.arange(16, dtype=torch.float32))
+        assert captured['publishes'] == 1
+        assert torch.equal(real_cpu(received), torch.arange(16, dtype=torch.float32))
 
     def test_isend_stages_payload_on_cpu(self, monkeypatch):
         from deepspeed.comm.torch import TorchBackend, StagedWork
-        captured = {}
+        captured = {'payload_reads': 0, 'payload_copy_backs': 0}
+        backend = TorchBackend.__new__(TorchBackend)
+        payload = torch.arange(16, dtype=torch.float32, device='mps')
+        real_to = torch.Tensor.to
+        real_copy = torch.Tensor.copy_
+
+        def spy_to(tensor, *args, **kwargs):
+            result = real_to(tensor, *args, **kwargs)
+            if tensor is payload and result.device.type == 'cpu':
+                captured['payload_reads'] += 1
+            return result
+
+        def spy_copy(tensor, source, *args, **kwargs):
+            result = real_copy(tensor, source, *args, **kwargs)
+            if tensor is payload:
+                captured['payload_copy_backs'] += 1
+            return result
 
         def fake_isend(tensor, dst=None, group=None, tag=0):
             captured['staged'] = tensor
             return FakeP2PWork()
 
+        monkeypatch.setattr(torch.Tensor, 'to', spy_to)
+        monkeypatch.setattr(torch.Tensor, 'copy_', spy_copy)
         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 captured['payload_reads'] == 1
         assert handle.wait() is True
+        assert captured['payload_reads'] == 1
+        assert captured['payload_copy_backs'] == 0
+
+    def test_irecv_failure_does_not_publish_uninitialized_staging(self, monkeypatch):
+        from deepspeed.comm.torch import TorchBackend
+        backend = TorchBackend.__new__(TorchBackend)
+        received = torch.full((16, ), 7.0, dtype=torch.float32, device='mps')
+        captured = {}
+
+        def fake_irecv(tensor, src=None, group=None, tag=0):
+            captured['staged'] = tensor
+            return FakeP2PWork(error=RuntimeError("receive failed"))
+
+        monkeypatch.setattr(torch.distributed, 'irecv', fake_irecv)
+        handle = backend.irecv(received, src=0)
+
+        assert captured['staged'].device.type == 'cpu'
+        with pytest.raises(RuntimeError, match="receive failed"):
+            handle.wait()
+        assert torch.equal(received, torch.full_like(received, 7.0))
+
+    def test_irecv_without_work_preserves_native_return_and_destination(self, monkeypatch):
+        from deepspeed.comm.torch import TorchBackend
+        backend = TorchBackend.__new__(TorchBackend)
+        received = torch.full((16, ), 7.0, dtype=torch.float32, device='mps')
+
+        monkeypatch.setattr(torch.distributed, 'irecv', lambda *args, **kwargs: None)
+        handle = backend.irecv(received, src=0)
+
+        assert handle is None
+        assert torch.equal(received, torch.full_like(received, 7.0))
 
 
 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':
+        # device_count(), which reports sockets on CPU and one physical device on MPS. Bypass
+        # the gate for both; MPS still needs spawned processes so each rank gets a Metal context.
+        device = get_accelerator().device_name()
+        if device not in ['cpu', 'mps']:
             return super()._launch_procs(num_procs, init_method)
         self.backend = 'gloo'
+        if device == 'mps':
+            self.non_daemonic_procs = True
+            self.reuse_dist_env = False
+            return self._launch_non_daemonic_procs(num_procs, init_method)
         torch.multiprocessing.set_start_method('forkserver', force=True)
         self._launch_daemonic_procs(num_procs, init_method)
 
@@ -215,6 +299,19 @@ def test(self):
             expected = torch.arange(length, dtype=torch.float32).to(device)
             assert torch.equal(received, expected)
 
+    def test_non_member_irecv_preserves_native_return(self):
+        rank = dist.get_rank()
+        device = get_accelerator().device_name()
+        subgroup = dist.new_group(ranks=[0])
+
+        if rank == 1:
+            received = torch.full((4096, ), 7.0, dtype=torch.float32, device=device)
+            handle = dist.irecv(received, src=0, group=subgroup)
+
+            assert handle is None
+            assert torch.equal(received, torch.full_like(received, 7.0))
+        dist.barrier()
+
 
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("num_elements", [128, 3])