[Core][Hardware][NVIDIA] Add custom all-reduce suspend hooks#47806
Draft
galletas1712 wants to merge 1 commit into
Draft
[Core][Hardware][NVIDIA] Add custom all-reduce suspend hooks#47806galletas1712 wants to merge 1 commit into
galletas1712 wants to merge 1 commit into
Conversation
Allow custom all-reduce to close all imported peer mappings before a process suspend and rebuild them from fresh IPC handles after resume. Captured CUDA graphs stay valid because kernels read peer pointers through stable device-side RankSignals/RankData tables that resume rewrites in place. Lifecycle failures are terminal (fail-fast, no rollback or cross-rank voting); repeated successful transitions are no-ops. Assisted-by: OpenAI Codex Assisted-by: Claude Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
galletas1712
force-pushed
the
schwinns/custom-allreduce-suspend-resume
branch
from
b6334af to
b46a96e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Add backend-owned suspend/resume hooks to
CustomAllreducefor checkpoint lifecycle support (part of #46877 / #47500; FlashInfer counterparts: flashinfer-ai/flashinfer#3727, flashinfer-ai/flashinfer#3745):prepare_for_suspend()collectively quiesces the TP group and closes all imported CUDA IPC peer mappings (wrapper-owned metadata/eager-staging mappings and native captured-graph mappings).reinit_after_resume()exchanges fresh IPC handles, reopens peer mappings, and rewrites the stable device-sideRankSignals/RankDatatables so previously captured CUDA graphs remain valid.Why indirection instead of stable-address remap: peer signal pointers were previously baked into captured kernel launches by value, and legacy
cudaIpcOpenMemHandlecannot remap at a requested address after resume. Kernels therefore now read peer signals through a device-side table that resume rewrites — the same pattern the existingRankDatatable already uses for graph-captured peer activations. A VA-stablecuMemMap-based design was considered and deferred: it requires FD-exchange infrastructure and cannot cover torch-allocated graph buffers, which force legacy IPC in any zero-copy design.Lifecycle semantics (modeled on flashinfer-ai/flashinfer#3727): fail-fast — a failed transition leaves the communicator in a terminal failed state and any use raises (no NCCL degradation during the suspend window); repeated successful calls are idempotent no-ops. No rollback or cross-rank voting.
Not a duplicate: no open PR implements custom-all-reduce IPC teardown/rebuild; #45611 only disables custom all-reduce under the cumem allocator and does not preserve mappings or captured graphs.
Test Plan
pytest tests/distributed/test_custom_all_reduce.pyon ≥2 CUDA GPUs (covers eager + captured-graph suspend/resume, idempotency, detached-use errors, close); 4 GPUs for thepp=2parametrization.uv pip install -e .) and ROCm compile check (suspend/resume itself is CUDA-only and raises on ROCm).benchmarks/kernels/benchmark_device_communicators.pybefore/after, to bound the cost of the signal-table indirection in the reduce kernels.Test Result
TODO — pre-commit hooks (ruff, clang-format, mypy, typos) pass; GPU runs and benchmark numbers pending. Post-resume validation is intentionally limited to cheap invariants (metadata counts, local signal/eager-buffer identity) with robust validation left as a
TODOinreinit_after_resume— layout comparison alone cannot prove mapping integrity.This PR was created with AI assistance. The submitting human must review every changed line and complete the GPU validation above before marking it ready.