mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
main
5
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 PyTorch MergeBot
|
2eb73612d3 |
Revert "Enable NCCLSymmetricMemory (RCCL) symmetric-memory support on ROCm (#192524)"
This reverts commit
|
||
|
Abhilash KaturuandChinmay Kuchinad
|
e64cfa6fe9 |
Enable NCCLSymmetricMemory (RCCL) symmetric-memory support on ROCm (#192524)
## Summary
This PR enables `NCCLSymmetricMemory` on ROCm. ROCm cannot reuse the CUDA path: HIP graph capture, RCCL window registration, and the device reduce APIs do not behave the same way. The changes below are the smallest set that make allocation, rendezvous, capture, and the three device ops work on the tested RCCL stack.
Almost all of the C++ under `torch/csrc` is `USE_ROCM`-guarded. Five changes are not, and do change CUDA behavior; they are listed under "Changes that affect the default path". `NCCL_HAS_SYMMEM_DEVICE_SUPPORT` on CUDA stays at NCCL 2.28.0.
## Problem
`NCCLSymmetricMemory` was CUDA-only. Opening the same gates on ROCm is not enough:
- RCCL `ncclMemAlloc` without CUMEM falls back to `hipMalloc`. Capture rejects that.
- A synchronous `hipMemset` or `hipMemcpy` during Global capture returns HIP 906 and kills the capture (901).
- An async memset recorded on the capturing stream becomes a graph node and can wipe a peer signal pad on every replay.
- RCCL 2.30.4 window registration fails during capture. 2.30.7 succeeds.
- A fused multi-destination reduce-scatter launch faults with `HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION` on gfx942 and gfx950.
- A retained handle after process-group destroy launched a barrier kernel through a dead communicator and hit `HSA_STATUS_ERROR_MEMORY_FAULT`.
## Solution
On ROCm, a per-device `NCCLDevCommManager` snapshots RCCL device-API support when a communicator registers. Without that support a handle is host-only. Fresh windows use CUMEM plus an async zero on a non-captured setup stream wherever RCCL supports capture-time allocation. Peer tables upload on the same stream, under a Relaxed capture-mode guard, so rendezvous returns with initialized device pointers. Device comms are created once per op and cached. Reduce-scatter launches one kernel per destination. `ProcessGroupNCCL` retires its registry entry through a hook that `NCCLComm` runs at the start of `abort()` and `destroy()`, and every ROCm accessor that can reach a communicator checks that its registration is still live. A handle whose communicator is gone raises; it is not rebuilt against a successor.
On ROCm the feature needs RCCL >= 2.30.4 and an `nccl_device.h` that compiles in a host translation unit, which CMake probes. Capture-time allocation additionally needs RCCL >= 2.30.7. There is no GPU- or ROCm-version-specific floor.
## Design choices
**`nccl_device.h` presence is not enough.** A host C++ TU cannot compile every RCCL that ships the header, so `cmake/External/rccl.cmake` try-compiles the installed header and defines `RCCL_DEVICE_HEADER_HOST_COMPATIBLE` when it passes. On ROCm, `NCCL_HAS_SYMMEM_DEVICE_SUPPORT` requires that probe and RCCL >= 2.30.4. The condition is written `!defined(USE_ROCM) || (...)`, so CUDA's threshold stays at NCCL 2.28.0. hipcc needs CUDA's extended-lambda macro and a local `__half operator+` because PyTorch defines `__HIP_NO_HALF_OPERATORS__`.
**Device API support is a runtime snapshot; without it the handle is host-only.** RCCL reports `deviceApiSupport=0` unless both `NCCL_CUMEM_ENABLE` and `NCCL_WIN_ENABLE` are set before communicator init. `register_comm` records `ncclCommQueryProperties` against that communicator identity, so setting the env vars later does not upgrade an existing comm. Rendezvous on such a comm still registers the window, so host-side collectives on the buffer work; the peer-pointer accessors and `barrier()` raise with a message naming the missing condition. `test_host_only_rendezvous_without_device_api_support` covers this.
**Capability, not capture state, selects the allocation path.** The CUMEM + setup-stream path is correct outside capture too, so keying on `capture_allocation_supported()` gives one path on a supported stack. HIP only reports capture state for the current stream, so this PR adds no process-wide capture query: the CUMEM path runs under the Relaxed guard, performs no capture-illegal host wait, and writes the pad and peer tables as host-side setup rather than graph nodes. The current-stream `TORCH_CHECK` on the other path is a diagnostic for stacks that cannot allocate under capture, not the safety mechanism.
**`TORCH_NCCL_SYMM_MEM_DISABLE_CAPTURE_ALLOC=1` keeps the plain path.** It is read inside `capture_allocation_supported()`, so it selects `ncclMemAlloc` plus synchronous memset for allocation *and* a synchronous `cudaMemcpy` for the rendezvous peer-table upload. That is the allocate-then-capture fallback for stacks that cannot allocate under capture, and it is how the two allocation paths were compared: both arms produce the same peer digest, which is what justifies keying on capability rather than on capture state.
**One shared setup stream per device.** Signal-pad zero and peer-table H2D need a stream that is not the capturing stream. It is created when the first comm with capture-allocation support registers, so every alloc and rendezvous reuses it rather than creating one per call. It is created regardless of the kill switch, and its use is serialized.
**A host-side wait is illegal under Global capture, even on a non-captured stream.** Moving the H2D copy off the capturing stream is not enough. `cudaStreamSynchronize` on the setup stream still returns HIP 900. `alloc()` therefore wraps its whole `ncclMemAlloc` / async-zero / synchronize sequence in `CUDAStreamCaptureModeGuard{Relaxed}`, and rendezvous does the same for the peer-table upload. The peer tables and signal pad are immutable for the object's lifetime, so they are initialized on the setup stream and are not recorded as graph nodes. `hipStreamWaitEvent` would also work, but it would only guarantee the tables are ready before later device work, not when rendezvous returns on the host.
**Fresh capture allocation requires RCCL >= 2.30.7.** Direct window registration during Global capture succeeds on 2.30.7 (mode stays Global, capture stays active). 2.30.4 fails in worker-side synchronization. The check is made at compile time and against `ncclGetVersion()` at registration. Older RCCL still works if the tensor is allocated and rendezvoused before capture.
**`ncclDevCommCreate` is a direct call, then cached per op.** First create during Global capture succeeds on 2.30.7, keeps Global mode and an active capture, and the captured graph replays. All-to-all and reduce-scatter each create a devcomm once and cache it under their own op key. There is no extra ROCm-only capture guard around create. The symmetric-memory barrier uses the signal pads directly and needs no devcomm.
**RCCL leaks Relaxed mode on window deregister and device-comm destroy.** Direct `ncclCommWindowDeregister` and `ncclDevCommDestroy` change the caller's capture thread mode from Global to Relaxed. Destroy during an active capture also invalidates the capture. Both calls sit in a balanced `CUDAStreamCaptureModeGuard{Relaxed}` so the caller gets its mode back.
**The manager destructor keeps `cudaDeviceSynchronize`.** HIP returns 900 for that wait under an active capture in both Global and Relaxed thread mode, so a Relaxed guard would not make teardown capture-safe. This enablement does not change destructor synchronization.
**Reduce-scatter is serialized per destination, not fused.** The fused launch faults on gfx942 and gfx950 with RCCL 2.30.7. Row tiles that write one destination are safe. ROCm drives the existing shared kernel with a one-slot `ReduceScatterOffsetsInfo` rather than getting a kernel of its own. With `n_owned == 1` the barrier index collapses to `blockIdx.x`, so the per-destination launches are expected to agree bit for bit with a fused launch of the same shape.
**`dst_block_size` stays a `uint16_t`.** The shared kernel layout is the CUDA struct and sizes above 65535 do not fit that field. This PR keeps the layout and validates the size on the host instead; see "Changes that affect the default path" for why that check is not ROCm-scoped.
**On ROCm, retirement runs inside `NCCLComm`, before RCCL sees the abort or destroy.** `initNCCLComm` registers the collective comm and installs a pre-invalidate hook on the `NCCLComm` as one step under the comm's own mutex, which `abort()` and `destroy()` also take, so an invalidation cannot land in between. The hook captures the group uid, the `ncclComm_t` (because `abort()` nulls the handle), and the registration generation, and `NCCLComm::destroy()` / `abort()` run it first (`NCCLUtils.cpp:362`, `:379`). Retirement only erases registry entries and makes no RCCL call, so it is safe on every abort path, including watchdog and wait-timeout aborts. The generation is what tells our entry apart from a successor that is handed the same `ncclComm_t` address. Because the hook covers every teardown path, the `~ProcessGroupNCCL` unregister loop is compiled out on ROCm (`ProcessGroupNCCL.cpp:1614`); on CUDA that loop is unchanged. nccl2 already retired before destroy and abort; it now also retires on a failed split or shrink (see below).
**Retirement does not call `ncclDevCommDestroy`.** It is collective (it deregisters the devcomm's resource window), so it cannot run on the unilateral abort path. The registry drops the entries, and comm teardown reclaims those windows: RCCL's `commFree` always runs `ncclDevrFinalize`, and NCCL's runs it for every comm with symmetric support, which `ncclDevCommCreate` requires. GIN contexts are not reclaimed that way, so this holds for single-node symmetric memory.
**Stale handles are rejected, not rebuilt.** Once a handle's registration is retired, every checked accessor raises `... is stale because its RCCL communicator was destroyed or replaced. Free the tensor and allocate a new one after initializing the successor process group.` Re-rendezvous is included: a cached handle (`NCCLSymmetricMemory.cu:876`) and another pointer into the same allocation (`:905`) both check liveness before reuse. An earlier revision rebuilt such a handle against a same-name successor. That rebuild starts a collective window registration based on rank-local registry state, and its test hung once in 13 runs in `destroy_process_group` on MI355X with RCCL 2.31.2. It did not reproduce in 264 runs on MI300X with RCCL 2.31.2 or in 6 CI runs, no stack was captured, and it was not root-caused, so this PR drops the rebuild. `test_stale_handle_rejected_across_same_name_restart` covers the replacement contract: the old handle raises from `barrier`, `put_signal`, `wait_signal`, and re-rendezvous, while a freshly allocated tensor on the successor works.
**Liveness checks are ROCm-only.** The use-after-destroy hazard is not ROCm-specific, but the check relies on every producer retiring before it invalidates its comm, which stock `ProcessGroupNCCL` does only on ROCm, so CUDA accessors are unchanged. Checked: the four peer-pointer accessors and `barrier()` (which also require device API support), `put_signal`, `wait_signal`, `get_window`, `has_multicast_support`, `get_group_name`, and the two CFT handle accessors. `get_group_name()` checks because resolving a communicator from a name is the hazard. `get_window_offset()`, `get_offset()`, `get_buffer_size()`, `get_rank()`, `get_world_size()`, and `get_device()` do not: they return immutable scalars, and callers that use the offset go through `get_window()`. The window destructor deregisters only if its exact registration, pointer and generation, is still live. These checks cover use after teardown, not use that races a concurrent abort.
**Host-side CFT is not enabled on ROCm.** `NCCL_HAS_HOST_CFT` gains `!defined(USE_ROCM)` because host-side CFT has not been validated on RCCL. `SymmMemCftHandleTest` skips on ROCm.
**Allocation size is rounded to `NCCL_WIN_REQUIRED_ALIGNMENT` on ROCm only.** An unrounded backing allocation produced wrong `put_wait_signal` data and a teardown timeout. Registration uses the same rounded extent. CUDA sizes are not changed.
## Changes that affect the default path
Five changes here are not under `USE_ROCM` and do change CUDA behavior. Listing them explicitly because the rest of the PR is ROCm-scoped and a reviewer should not have to grep for the exceptions.
- **`initNCCLComm` publishes only the collective comm.** Every comm `initNCCLComm` created was registered under `getGroupUid()`, last write wins, so a lazily created 2-rank send/recv comm could replace the group's collective comm, and symmetric memory on that group would then resolve to a comm that does not span the group. Registration now runs only `if (!singleP2POp)` (`ProcessGroupNCCL.cpp:3326`). On CUDA, a group that has only ever created a P2P comm therefore has no registry entry, and rendezvous on it raises `get_comm`'s "host communicator ... not found" error instead of using the P2P comm. No dedicated test.
- **Reduce-scatter silently truncated blocks larger than 65535.** `ReduceScatterOffsetsInfo::dst_block_size` is a `uint16_t` and the host wrote `static_cast<uint16_t>(owned_sizes[j])` with no bound; the only prior check was on block *count*. A taller shard reduced `owned_sizes[j] % 65536` rows and returned wrong data with no error. Now a host-side `TORCH_CHECK` (`nccl_reduce_scatter_offset.cu:345`), placed before the ROCm/CUDA launch split so both paths get it. The inputs this rejects are exactly the ones that were already producing wrong results. `test_reduce_scatter_offset_rejects_oversized_block` asserts both that the first oversized block raises and that the largest block that still fits reduces correctly.
- **Cached device comms survived a communicator replacement.** `register_comm` is last-write-wins, but a cached `ncclDevComm` is built from one specific host comm. When a successor process group registered a different comm under the same uid, `get_devcomm` kept handing kernels devcomms naming a comm the registry no longer owned, and `~NCCLDevCommManager` would pass the successor's comm to `ncclDevCommDestroy` alongside a predecessor-built devcomm. `register_comm` now evicts that group's devcomms on replacement (`nccl_devcomm_manager.hpp:261`); re-publishing the same pointer is not a replacement. Identity is the pointer, so on CUDA, where `~ProcessGroupNCCL` keeps the entry of a destroyed comm, a successor at a recycled `ncclComm_t` address still skips the eviction; nccl2 and ROCm retire before invalidating, which rules that out. Covered by `NCCLDevCommManagerTest.ReplacementRegistrationEvictsDevComms`, which also asserts that a bystander group's devcomm survives and that a predecessor's late retirement cannot unpublish its successor.
- **The one-sided ops dereferenced an unchecked `dynamic_cast`.** `nccl_put_signal` and `nccl_wait_signal` are schema'd over the `SymmetricMemory` base class, so a handle from another backend is a legal argument, and both used the cast result without testing it. Now `TORCH_CHECK`ed (`nccl_extension.cu:375`, `:410`), matching the sibling ops. Covered by `NCCLOneSidedOpHandleTypeTest::test_one_sided_ops_reject_foreign_handle`.
- **nccl2 left the parent registered when split or shrink failed.** On failure both cleared `nccl_comm_` without retiring, leaving the registry pointing at a parent the failed call may have aborted. Both now call `retireComm()` first (`nccl2/ProcessGroupNCCL.cpp:399`, `nccl2/ShrinkNCCL.cpp:100`). No test: it needs a split or shrink that fails.
Also visible on CUDA, but not behavioral:
- The two ops copy the cached `ncclDevComm` instead of binding a reference into the registry (`nccl_all_to_all_nd.cu:183`, `nccl_reduce_scatter_offset.cu:221`), since eviction can now erase the entry. `get_devcomm`'s signature is unchanged; its comment now says the reference dangles after replacement.
- `NCCLComm` gains `preInvalidateHook_` on every platform, so the class layout does not depend on `USE_ROCM`. It is only set and run on ROCm.
- `nccl_dev_cap.hpp` only adds ROCm terms, so the CUDA truth values of `NCCL_HAS_SYMMEM_DEVICE_SUPPORT` and `NCCL_HAS_HOST_CFT` are unchanged. The `NCCL_HAS_DEVCOMM` / `NCCL_HAS_ONE_SIDED_API` hunk merges two byte-identical `#if` conditions.
- `#include <vector_types.h>` is dropped from `NCCLSymmetricMemory.cu`; nothing there uses a CUDA vector type, and the header does not exist on ROCm.
- `_is_nccl_symmem_available()` is new and only used by the test suite to pick a skip predicate. The rest of the `init.cpp` hunk is clang-format reflow of existing lambdas.
## Tests
`test/distributed/test_nccl.py` no longer carries a blanket `TEST_WITH_ROCM` skip on the symmetric-memory classes. It is replaced by narrower gates, all no-ops on CUDA:
- `skip_if_rccl_symmem_not_compiled`: the build-time result, exposed to Python as `_is_nccl_symmem_available()`.
- `skip_if_rccl_lt_2_30_4`: the RCCL host device-API floor.
- `skip_if_rccl_lt_2_31_2`: `test_put_wait_signal` and `test_nccl_symmem_signal_rank_out_of_bounds`, which are validated only on RCCL 2.31.2.
- Tests that allocate during capture skip on RCCL < 2.30.7.
New coverage:
- `NCCLSymmetricMemoryLifecycleTest` (ROCm-only, 8 items, each on `nccl` and `nccl-legacy`; only `nccl-legacy` goes through stock `ProcessGroupNCCL`): a retained handle is rejected after abort, after a wait-timeout abort, and after a watchdog abort, and across a same-name restart as described above.
- `NCCLSymmetricMemoryCapabilityGateTest` (ROCm-only, 3 items): the host-only path without device API support, capture allocation rejected when CUMEM/WIN are enabled after init, and the kill switch.
- `NCCLSymmetricMemorySubgroupTest` (ROCm-only, 1 item): subgroup-first rendezvous in a fresh process.
- `NCCLSymmetricMemoryRestartTest` (**not** ROCm-gated, 2 items): allocation after a default-PG restart, and a successor PG surviving its predecessor's teardown.
- `NCCLOneSidedOpHandleTypeTest` (**not** ROCm-gated, 1 item): one-sided ops reject a non-NCCL handle.
- `NCCLSymmetricMemoryTest::test_reduce_scatter_offset_rejects_oversized_block` (**not** ROCm-gated), plus the ROCm-only capture cases `test_nccl_symmem_capture_allocation_kill_switch` and `NCCLSymmetricMemoryNccl2Test::test_nccl_symmem_fresh_cuda_graph`.
- `NCCLDevCommManagerTest` (C++): `ReplacementRegistrationEvictsDevComms` on both platforms, and three ROCm-only generation-token cases.
Two CI changes, both to shared config:
- `.ci/pytorch/test.sh:590`: the H100 `sm90-symm` filter widens from `NCCLSymmetricMemoryTest` to `"NCCLSymmetricMemoryTest or NCCLOneSidedOpHandleTypeTest or NCCLSymmetricMemoryRestartTest"`, so the CUDA-visible fixes run on CUDA hardware with P2P.
- `.ci/pytorch/test.sh:1907`: `cpp/NCCLDevCommManagerTest` runs in the distributed C++ block shared by CUDA and ROCm on shard 1. On a ROCm build without device support the binary contains one `GTEST_SKIP` case. This branch is rebased onto #199440; before it, every ROCm C++ test in that block collected 0 items and still reported success.
Still skipped on ROCm: `test_nccl_symmem_rendezvous_subgroup` (RCCL hangs in `ncclRmaCeInit`'s bootstrap barrier when WORLD already holds a window) and `SymmMemCftHandleTest` (host-side CFT is not enabled on ROCm).
## Validation
### CI on this head
Read from the uploaded `test-reports` XML and the per-test C++ logs on [`f07e71b`](https://github.com/pytorch/pytorch/commit/f07e71ba3f197b2814383ce1dcf39447d4007202).
| Job | GPU | `test_nccl.py` symmetric-memory items | `cpp/NCCLDevCommManagerTest` |
|---|---|---|---|
| [rocm-preview](https://github.com/pytorch/pytorch/actions/runs/37410391165) `linux-noble-rocm-preview-py3.12-mi350` | gfx950 | **51 passed, 5 skipped** | [**11 passed**](https://github.com/pytorch/pytorch/actions/runs/37410391165/job/112119195812) |
| [trunk](https://github.com/pytorch/pytorch/actions/runs/37410384883) `linux-noble-rocm-py3.11-mi350` | gfx950 | 56 skipped | [1 skipped](https://github.com/pytorch/pytorch/actions/runs/37410384883/job/112125954872) |
| [periodic-rocm-mi300](https://github.com/pytorch/pytorch/actions/runs/37410383794) `linux-noble-rocm-py3.11-mi300` | gfx942 | 56 skipped | [1 skipped](https://github.com/pytorch/pytorch/actions/runs/37410383794/job/112125941200) |
| [periodic-rocm-mi200](https://github.com/pytorch/pytorch/actions/runs/37410386548) `linux-jammy-rocm-py3.11-mi200` | gfx90a | 56 skipped | [1 skipped](https://github.com/pytorch/pytorch/actions/runs/37410386548/job/112128613809) |
| [trunk](https://github.com/pytorch/pytorch/actions/runs/37410384883) `linux-jammy-cuda13.2-py3.11-gcc11` | T4 | 56 skipped (no P2P) | [8 passed](https://github.com/pytorch/pytorch/actions/runs/37410384883/job/112110745550) |
| [`sm90-symm`](https://github.com/pytorch/pytorch/actions/runs/37410384954/job/112111113881) CUDA 13.2 | H100 | 36 passed, 2 skipped | not run |
Preview's 5 skips are the two signal tests (CI ships RCCL 2.30.7), the two `SymmMemCftHandleTest` cases, and the subgroup test. All 8 `NCCLSymmetricMemoryLifecycleTest` items pass there, including the same-name restart on both backends. The other three ROCm labels skip 51 items on the build-time gate (their RCCL is older than 2.30.4 or its `nccl_device.h` fails the host-compile probe), plus 3 on RCCL < 2.30.7 and 2 on CFT, and their gtest binary holds only the `GTEST_SKIP` case. Preview is therefore the only ROCm label with live coverage. Widening that is a matter of the trunk and periodic images carrying a qualifying RCCL, not of test gating.
On H100 the `sm90-symm` filter selects 38 items. The oversized-block, foreign-handle, and both restart cases are among the 36 that pass; the two skips are `test_multicast_ptr` (`NCCL_NVLS_ENABLE=0` on that runner) and the ROCm-only `test_nccl_symmem_capture_allocation_kill_switch`. On T4 the gtest runs the 7 existing cases and `ReplacementRegistrationEvictsDevComms`; the three generation-token cases are ROCm-only.
### Locally
8x MI300X (gfx942), HIP 7.17.26383, RCCL 2.31.2, via `run_test.py` as CI invokes it, on the same patch before the rebase:
- `distributed/test_nccl.py`: all 56 symmetric-memory items collected, 53 passed and 3 skipped (subgroup and the two CFT cases). The signal tests run here because RCCL is 2.31.2.
- `test_stale_handle_rejected_across_same_name_restart`: 100 runs across four GPU pairs, both backends per run, no failure or hang.
- `cpp/NCCLDevCommManagerTest`: 11 passed.
- `lintrunner` clean on the changed files.
### Measurements CI does not reproduce
The two-arm allocation digest comparison, the in-capture alloc/rendezvous/replay probe, and the per-config reduce-scatter kernel comparison were run on an 8x MI355X box (HIP 7.16.26344, RCCL 2.30.7, `torch.version.rocm` 10.1.0). They are recorded in the evidence notes and available on request.
Fixes #168480
Authored with the assistance of an AI coding assistant.
Co-authored by @chinmaydk99
Pull Request resolved: https://github.com/pytorch/pytorch/pull/192524
Approved by: https://github.com/jeffdaily
Co-authored-by: Chinmay Kuchinad <ChinmayDattanand.Kuchinad@amd.com>
|
||
|
Ke Wen
|
912ffb022e |
Add C++ tests for NCCLDevCommManager (#177473)
This commit adds comprehensive C++ tests for the NCCLDevCommManager class,
which manages NCCL device communicators for symmetric memory operations.
We mainly want to test if the manager can help developer cache the device communicator:
```
void foo(group) {
auto devcomm_opt = get_devcomm(group_name);
if (!devcomm_opt) {
ncclDevComm devcomm = ncclDevCommCreate(...);
// Cached under group_name and function name "foo"
devcomm_opt = register_devcomm(group_name, devcomm);
}
}
```
Changes:
- Add test/cpp/c10d/NCCLDevCommManagerTest.cpp with tests covering:
* register_comm and get_comm functionality
* register_devcomm and get_devcomm with default and custom keys
* Error handling for duplicate registrations and non-existent keys
* Multiple groups and string_view parameter support
- Add NCCLDevCommManagerTest to CMakeLists.txt with proper compile
definitions (USE_CUDA, USE_NCCL) to enable NCCL symmetric memory device
support detection
- Exclude test/cpp/**/*.cpp from RAWTHROW linter since test files use
EXPECT_THROW which is a gtest macro
The tests require CUDA and NCCL version >= 2.28.0 with symmetric memory
device support. Tests will skip gracefully if these requirements are not met.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/177473
Approved by: https://github.com/Skylion007
ghstack dependencies: #177472, #177459, #177463, #177380
|
||
|
PyTorch MergeBot
|
cc57e0e7ca |
Revert "Add C++ tests for NCCLDevCommManager (#177473)"
This reverts commit
|
||
|
Ke Wen
|
0eeab58118 |
Add C++ tests for NCCLDevCommManager (#177473)
This commit adds comprehensive C++ tests for the NCCLDevCommManager class,
which manages NCCL device communicators for symmetric memory operations.
We mainly want to test if the manager can help developer cache the device communicator:
```
void foo(group) {
auto devcomm_opt = get_devcomm(group_name);
if (!devcomm_opt) {
ncclDevComm devcomm = ncclDevCommCreate(...);
// Cached under group_name and function name "foo"
devcomm_opt = register_devcomm(group_name, devcomm);
}
}
```
Changes:
- Add test/cpp/c10d/NCCLDevCommManagerTest.cpp with tests covering:
* register_comm and get_comm functionality
* register_devcomm and get_devcomm with default and custom keys
* Error handling for duplicate registrations and non-existent keys
* Multiple groups and string_view parameter support
- Add NCCLDevCommManagerTest to CMakeLists.txt with proper compile
definitions (USE_CUDA, USE_NCCL) to enable NCCL symmetric memory device
support detection
- Exclude test/cpp/**/*.cpp from RAWTHROW linter since test files use
EXPECT_THROW which is a gtest macro
The tests require CUDA and NCCL version >= 2.28.0 with symmetric memory
device support. Tests will skip gracefully if these requirements are not met.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/177473
Approved by: https://github.com/Skylion007
ghstack dependencies: #177472, #177459, #177463, #177380
|