Route isend/irecv to nonblocking backend methods and stage them as async on MPS

[PKUWZP]

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, DeepSpeed Comm. Backend v1 #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 isends, rank 1 irecvs, 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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f555e9ede9

ℹ️ 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".

[FU-max-boop]

Thanks for the concise follow-up and for the credit. I checked the CI execution after 5be039e and found one remaining coverage gap: the CPU job collected TestDistIsendIrecv::test, but marked it SKIPPED (job log).

DistributedTest._launch_procs() compares get_accelerator().device_count() with world_size; the CPU runner reports one socket, so world_size = 2 never reaches real Gloo/Work there. The test can still exercise NCCL on a multi-GPU runner, but it does not currently provide the requested real two-rank Gloo regression.

A minimal way to preserve both paths is to bypass only that per-device gate on CPU, following the existing TestRealCheckpointUniversalConversionTPxPP precedent:

    def _launch_procs(self, num_procs, init_method):
        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)
        return self._launch_daemonic_procs(num_procs, init_method)

Setting the CPU branch explicitly to Gloo avoids silently selecting oneCCL when its bindings are present. An assert dist.get_backend() == "gloo" in that branch would also make the intended coverage observable. The same test then runs CPU/Gloo in the CPU job and the accelerator-native backend in GPU CI. Happy to provide a patch/cherry-pick if useful.

[PKUWZP]

@FU-max-boop Thanks for the comments, good catch! — confirmed from the job log that the CPU runner collected but skipped the two-rank test (one socket). Applied your suggestion in 1504ea6: TestDistIsendIrecv bypasses the per-device gate on CPU and pins backend = 'gloo', following the TestRealCheckpointUniversalConversionTPxPP precedent.

One extra fix surfaced by the bypass: ShareMemCommBuilder claimed compatibility off Linux, so any CPU-backend init_distributed on macOS died JIT-building the Linux-only shm kernels — it now reports incompatible there and build_shm_op degrades to None. Verified the two-rank gloo run passes on a single-socket machine (real Work handles end to end), and the MPS suite still passes.

[delock]

why there is adam change? Should it be moved to a separate PR if need to fix a different issue?

[PKUWZP]

Fair question — it's here because of CI sequencing rather than scope: after master was merged into this branch, the cpu-torch-latest job re-ran and failed on test_fused_adam_matches_torch (added in #8293), which blocked this PR. Digging into that failure exposed a latent bug in this file: multi_tensor_adam created its Adam_Optimizer once with default arguments and ignored the mode parameter, so FusedAdam(adam_w_mode=False) on the CPU backend always applied decoupled (AdamW) weight decay. The commit here is the minimal fix (one optimizer instance per mode) plus trimming the test's bf16 rows, whose torch-in-bf16 reference was never valid.

Happy to split it into its own PR if you prefer — the sequence would be: merge the kernel fix first, then update this branch so cpu-torch-latest stays green here. Just say the word and I'll extract it.

[PKUWZP]

Done — extracted to #8307 and reverted here (1493124f7), so this PR is back to P2P-only scope. Note: cpu-torch-latest on this PR will fail on test_fused_adam_matches_torch[fp32-adam] until #8307 merges; after it lands, a branch update here turns it green again.

[delock]

Hi @PKUWZP I'm generally okay with this PR. The only concern is adam behavior change should be submitted in a seperate PR, because it is not related to comms.

[PKUWZP]

Heads-up for reviewers: the CPU Adam mode fix was split out to #8307 per review feedback. Until #8307 merges, this PR's cpu-torch-latest check fails on test_fused_adam_matches_torch[fp32-adam] — a latent kernel bug unrelated to this PR's P2P changes. Merge order: #8307 first, then update this branch.

[FU-max-boop]

Thanks again for carrying the narrower patch. This is explicitly not a request to expand #8303; the current PR should stay focused, and the contract below would be handled only as a separate follow-up after this PR settles.

Would this minimal contract match maintainer expectations for a CPU-staged P2P Work?

1. Expose is_completed() as wrapper-level completion rather than blindly returning the underlying flag. If the underlying request is pending, return False. Once it is terminal, confirm its outcome through a path that cannot wait on still-in-flight work: on success, publish the staged receive buffer exactly once before returning True; on failure, do not publish it and preserve the failure for wait() to surface. A successful receive must never report completed while its copy-back is still pending.
2. Forward only the *args, **kwargs accepted by the wrapped wait() and preserve its return or exception; the wrapper would not add or normalize portable timeout support. Finalize only after backend-confirmed success. A timeout or other non-success must not publish receive data, and staging storage must remain alive while the underlying work may still reference it, without promising that a timed-out operation is retryable.
3. Make staging directional:
   • isend is input-only: copy MPS -> CPU once, retain that CPU payload through completion, and never copy it back.
   • irecv is output-only: allocate the CPU receive target without copying old MPS contents, then copy CPU -> MPS once and only after confirmed success.
4. Keep terminal state and successful receive finalization idempotent across repeated wait() / is_completed() calls.

I would deliberately exclude get_future() and result() from this P2P contract: PyTorch explicitly excludes Gloo/MPI peer-to-peer operations from get_future(), and Gloo SendWork / RecvWork do not implement result(). The public P2P surface does include is_completed() and wait(), while the current Gloo failure path can mark completion before rethrowing.

If this direction looks right, I can prepare that separate follow-up with behavior tests for success, error/no-copy, poll-before-wait, repeated completion, and both staging directions, plus an MPS copy-count and copy-cost baseline.

[PKUWZP]

Hi @PKUWZP I'm generally okay with this PR. The only concern is adam behavior change should be submitted in a seperate PR, because it is not related to comms.

@delock I opened the #8307 to fix the adam behavior changes in a separate PR.

[delock]

Thanks again for carrying the narrower patch. This is explicitly not a request to expand #8303; the current PR should stay focused, and the contract below would be handled only as a separate follow-up after this PR settles.

Would this minimal contract match maintainer expectations for a CPU-staged P2P Work?

1. Expose is_completed() as wrapper-level completion rather than blindly returning the underlying flag. If the underlying request is pending, return False. Once it is terminal, confirm its outcome through a path that cannot wait on still-in-flight work: on success, publish the staged receive buffer exactly once before returning True; on failure, do not publish it and preserve the failure for wait() to surface. A successful receive must never report completed while its copy-back is still pending.

2. Forward only the *args, **kwargs accepted by the wrapped wait() and preserve its return or exception; the wrapper would not add or normalize portable timeout support. Finalize only after backend-confirmed success. A timeout or other non-success must not publish receive data, and staging storage must remain alive while the underlying work may still reference it, without promising that a timed-out operation is retryable.

3. Make staging directional:

   • isend is input-only: copy MPS -> CPU once, retain that CPU payload through completion, and never copy it back.
   • irecv is output-only: allocate the CPU receive target without copying old MPS contents, then copy CPU -> MPS once and only after confirmed success.

4. Keep terminal state and successful receive finalization idempotent across repeated wait() / is_completed() calls.

I would deliberately exclude get_future() and result() from this P2P contract: PyTorch explicitly excludes Gloo/MPI peer-to-peer operations from get_future(), and Gloo SendWork / RecvWork do not implement result(). The public P2P surface does include is_completed() and wait(), while the current Gloo failure path can mark completion before rethrowing.

If this direction looks right, I can prepare that separate follow-up with behavior tests for success, error/no-copy, poll-before-wait, repeated completion, and both staging directions, plus an MPS copy-count and copy-cost baseline.

Hi @FU-max-boop, thanks for the thorough write-up.

I think directional staging (point 3) is the one worth a follow-up now. It's not hypothetical — with the current in-out staging, every isend does a redundant CPU→MPS copy-back and every irecv does a redundant MPS→CPU pre-copy. On the pipeline P2P path those are real, per-call costs. Please go ahead with that patch.

One thing to keep in mind while doing it: once irecv allocates with empty_like instead of pre-copying, a copy-back on a failed transfer would publish uninitialized memory into the user's tensor. Today's code happens to get this right (an exception from work.wait() skips copy_back()), so the follow-up mainly needs to not regress it — worth a test.

For is_completed() and wait(timeout=...): no caller in DeepSpeed uses either today (all isend/irecv call sites just do a bare .wait()), so we'd better harden them when we actually adopt those features. In the current implementation both fail loudly on use — AttributeError and TypeError respectively — which is a perfectly good outcome for an unused path. A short comment on StagedWork noting that this is deliberate, and that the signature must not be widened to *args/**kwargs (which would turn a loud failure into a silent one), should be enough for now.

Agreed on excluding get_future()/result() — matching what real gloo SendWork/RecvWork expose is the right bar.