Fix nonblocking P2P staging on MPS - #8301
Conversation
Signed-off-by: Fu Xiaonan <ht3fudatou@163.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c495d624b2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if result is False: | ||
| self._discard() |
There was a problem hiding this comment.
When the underlying Work.wait(timeout=...) returns False because the operation is still pending, _discard() permanently removes the completion callback and releases the staged CPU buffers. A later successful wait can therefore neither safely finish the transfer nor copy an irecv result back to the MPS tensor, causing lost receive data and potentially dropping an in-flight send buffer; retain the staging state on a timeout and only discard it after a confirmed terminal failure.
Useful? React with 👍 / 👎.
| 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) |
There was a problem hiding this comment.
Commit b7ee546299de3a8ade46b1f7d5ee5ac3d1c0d3bd is a non-merge commit but its message has no Signed-off-by trailer, so it does not meet this repository's commit requirements; recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
PKUWZP
left a comment
PKUWZP
left a comment
There was a problem hiding this comment.
Thanks so much for the PR, a couple of things here:
-
The
comm.pyrerouting changesisend/irecvsemantics for every backend (NCCL included) — it revives previously-dead async paths in the 1-bit comm modules. It's a global behavior change riding in an "MPS staging" PR; it should be called out prominently and watched on the CUDA CI, not just macOS. -
StagedWorkgrew locks, a weakref identity map, future wrapping, and__getattr__delegation (~100 lines) to preserveresult()/get_future_result()semantics that no current DeepSpeed caller uses. I feel it add unnecessary complexity — callers only use.wait()and occasionallyis_completed(). Can we drop the future/result machinery until something needs it.
|
@FU-max-boop Also I found that tests are all based on FakeTensor, FakeWork, monkeypatched staging predicate etc. — they elaborate implementation details (copy counts, weakref retention) rather than behavior, and none of them tests real torch.distributed. A 2-rank gloo test with real Work objects (like the one I ran) would be stronger and could run on any CI, not just macOS. |
|
Thanks @PKUWZP for the detailed review and for carrying the narrower implementation in #8303. I agree with both review points: #8301's locks/weakrefs/future/result machinery is larger than the current call surface justifies, and its fake-heavy tests do not replace a real two-rank #8303 now carries the core routing change, focused MPS staging coverage, and a real two-rank Gloo path, including the CPU skip fix in This is coordination, not a claim that #8303 has merged. I'll leave Thanks again for the concrete feedback. |
…ync on MPS (deepspeedai#8303) ## Summary Resolves @delock's review note on deepspeedai#8293 (deepspeedai#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, deepspeedai#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 deepspeedai#8301 deepspeedai#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>
2 participants
Summary
deepspeed.comm.isendandirecvthrough the backend's nonblocking methodsWorkcompletesWorkpolling, result, future, and error semantics while restoring original tensor identitiesThis follows the unresolved MPS review note on #8293: #8293 (comment). The existing wrapper only detects collectives with an
async_opargument, butisendandirecvare asynchronous by contract andirecvhas no such parameter. In addition, the public wrappers currently dispatch them to blocking backend methods.Validation
pytest tests/unit/comm/test_p2p.py— 16 passedpytest --forked tests/unit/comm/test_p2p.py— 16 passedpre-commit run --files deepspeed/comm/comm.py deepspeed/comm/torch.py tests/unit/comm/test_p2p.pyThe tests use deterministic fake MPS tensors and work handles to cover wait, polling, result and future completion, failure paths, no-grad copy-back, exactly-once behavior, and staging lifetime.