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](https://github.com/deepspeedai/DeepSpeed/pull/8303#issuecomment-5393623333).
## 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>
