TorchComms Integration for MSCCL++

[michael-beebe]

What This PR Does

This PR adds TorchComms support to MSCCL++, allowing PyTorch users to use MSCCL++ collectives through the TorchComms API with a single line:

comm = torchcomms.new_comm("mscclpp", device, name="grad_sync")
comm.all_reduce(tensor, torchcomms.ReduceOp.SUM, False)

This is valuable because it gives PyTorch training frameworks (torchtitan, FSDP2, etc.) a clean way to use MSCCL++ for high-performance collectives without LD_PRELOAD hacks or custom CUDA kernel code. Users can run MSCCL++ for the hot-path collectives (allreduce, allgather) and NCCL for everything else — mixed-backend training with no code changes.

Architecture

Communicator Lifecycle

When a user calls torchcomms.new_comm("mscclpp", device), TorchComms dlopen's our _comms_mscclpp.*.so module and calls init(), which:

1. Bootstrap — discovers rank/world_size from torchrun environment, exchanges a UniqueId through c10d::Store (rank 0 generates, others read), creates the MSCCL++ Communicator with a TcpBootstrap
2. Scratch buffer — allocates 128MB via GpuBuffer (cuMemMap) for native algorithms that need intermediate storage
3. Executor — creates the DSL plan executor (used by DSL algorithms, ignored by native ones)
4. Algorithm collection — calls AlgorithmCollectionBuilder::buildDefaultAlgorithms() which registers 12 native algorithms + 2 DSL plans, then wires up the topology-aware algorithm selector
5. Event pool — pre-allocates a pool of 256 reusable CUDA events for async work tracking

What Happens When You Call a Collective

comm.all_reduce(tensor, torchcomms.ReduceOp.SUM, False)
    │
    ▼
TorchCommMSCCLPP::all_reduce()
    │  validates reduce op (SUM/MIN only)
    │  ensures tensor is contiguous
    │
    ▼
TorchCommMSCCLPP::executeCollective("allreduce", sendbuf, recvbuf, size, dtype, ...)
    │
    │  1. Builds a CollectiveRequest with world_size, nRanksPerNode,
    │     rank, buffer pointers, message size, stream, dtype
    │
    │  2. Calls algorithmCollection_.selectAlgorithm(request)
    │     → selector considers message size, NVLS support, compute
    │       capability, symmetric memory, CUDA graph capture mode
    │     → returns the best algorithm (e.g., nvls_warp_pipeline for 4MB)
    │
    │  3. Creates TorchWorkMSCCLPP handle, records start GPU event
    │
    │  4. Calls algo->execute(comm, input, output, size, dtype, op, stream, executor)
    │     → native algorithms launch a CUDA kernel directly
    │     → DSL algorithms use the executor to interpret a JSON plan
    │
    │  5. Records end GPU event, returns the work handle
    │
    ▼
TorchWorkMSCCLPP (returned to caller)
    │  wait() → cudaStreamWaitEvent on caller's stream (GPU-side, no CPU block)
    │  checkStatus() → polls GPU events for completion/timeout

Component Diagram

torchcomms.new_comm("mscclpp", device)
    │
    ▼
TorchCommMSCCLPPPy.cpp          ← pybind11 module + dynamic loader interface
    │
    ▼
TorchCommMSCCLPP.cpp/hpp        ← backend class (init, finalize, collective dispatch)
    │
    ├── TorchCommMSCCLPPBootstrap.cpp/hpp  ← rank discovery via c10d::Store
    ├── TorchWorkMSCCLPP.cpp/hpp           ← GPU event pool + async work tracking
    │
    ▼
AlgorithmCollection::selectAlgorithm()   ← MSCCL++ native algorithm selection
    │
    ▼
Algorithm::execute()                      ← GPU kernel launch (native or DSL)

The backend is a thin adapter. It does not implement any collective algorithms — it delegates entirely to MSCCL++'s AlgorithmCollection, which selects the optimal native algorithm based on message size, topology, NVLS support, and compute capability.

Files to Review

Core Backend (focus here)

File	Lines	What It Does
TorchCommMSCCLPP.hpp	~180	Backend class declaration. All public methods, private members with comments.
TorchCommMSCCLPP.cpp	~510	The main implementation. init() bootstraps and builds the AlgorithmCollection. executeCollective() is the central dispatch — builds a CollectiveRequest, calls selectAlgorithm(), executes. Unsupported ops throw with NCCL/RCCL guidance.
TorchCommMSCCLPPBootstrap.hpp/cpp	~130	Rank discovery. Rank 0 generates a UniqueId, writes to c10d::Store, other ranks read it. Same pattern as TorchComms' NCCL backend.
TorchWorkMSCCLPP.hpp/cpp	~230	GPU event pool (amortizes cudaEventCreate cost) + async work handle. wait() uses cudaStreamWaitEvent for GPU-side sync — no CPU blocking.
TorchCommMSCCLPPPy.cpp	~60	Minimal pybind11 module. Exposes the class + DynamicLoaderInterface for TorchComms' dlopen discovery.
CMakeLists.txt (torchcomm)	~115	FetchContent for torchcomms headers, links mscclpp + PyTorch + pybind11.

Build System

File	Change
CMakeLists.txt (root)	2 lines: adds MSCCLPP_BUILD_EXT_TORCHCOMMS option (OFF by default) and add_subdirectory()

Tests, Benchmarks, Docs

Tests (6 files, ~960 lines), benchmarks (3 files, ~500 lines), and docs (quickstart.md) are straightforward and lower review priority.

Supported Collectives

Collective	Native Algorithms	Notes
AllReduce	allpair_packet, nvls_packet, packet, nvls, nvls_warp_pipeline, nvls_block_pipeline, fullmesh, rsag, rsag_pipeline, rsag_zero_copy	Auto-selected by message size + topology. SUM and MIN reduction ops.
AllGather	fullmesh, fullmesh2	Auto-selected by message size.
ReduceScatter	(none natively)	Dispatched if a DSL algorithm is registered.
AllToAll	(none natively)	Dispatched if a DSL algorithm is registered.
Broadcast, Reduce, Send/Recv, Barrier, Scatter, Gather	Not supported	Each throws with explicit message suggesting NCCL/RCCL.

Key Design Decisions

1. Thin adapter, not a reimplementation.
The backend calls AlgorithmCollection::selectAlgorithm() and Algorithm::execute(). It does not contain any collective kernel code. Algorithm registration, selection logic, and kernel implementations all live in MSCCL++ core.

2. Same algorithm selector as the NCCL extension.
We reuse algorithm_selector.hpp from src/ext/nccl/ so the TorchComms path selects the same algorithms as the LD_PRELOAD NCCL shim. This avoids divergence and ensures consistent behavior.

3. Shared library linking (not static).
The module links against libmscclpp.so (not mscclpp_static.a) to avoid dual-singleton crashes. mscclpp_collectives.so links against the shared lib, so if we statically linked, there would be two copies of singletons like UnixSocketServer::instance().

4. GpuBuffer for scratch allocation.
Scratch memory is allocated via mscclpp::GpuBuffer (cuMemMap) instead of plain cudaMalloc. This registers POSIX file descriptors in the unix socket server, which is required for cross-rank IPC sharing. Plain cudaMalloc causes "Requested fd not found" crashes.

5. Build-gated behind MSCCLPP_BUILD_EXT_TORCHCOMMS=OFF.
No impact on existing builds. TorchComms headers are fetched on-demand via CMake FetchContent only when the option is enabled.

6. GPU event pooling.
Every collective call needs 2 CUDA events (start + end) for async tracking. Creating/destroying events costs ~5-10μs each. The pool amortizes this across thousands of collective calls in a training loop.

7. User-defined algorithms via AlgorithmCollectionBuilder singleton.
Custom algorithms (DSL or native) are configured on the builder before creating the TorchComms communicator. The backend picks them up during init(). No algorithm registration API lives on the backend itself.

Limitations

• Only single-tensor collective variants are implemented. MSCCL++'s Algorithm::execute() operates on contiguous buffers (one input pointer, one output pointer), so the backend implements all_gather_single and reduce_scatter_single but not the tensor-list variants. The tensor-list variants throw with guidance to use the single-tensor variant instead.
• Unsupported collectives throw at runtime. Broadcast, reduce, send/recv, barrier, scatter, and gather throw a RuntimeError with an explicit message naming the operation and suggesting the caller use a separate NCCL/RCCL communicator. This is the expected pattern for mixed-backend training.
• Algorithm selector is duplicated from the NCCL extension. The backend includes algorithm_selector.hpp from src/ext/nccl/ rather than sharing it through a common path. A TODO in the code notes this should be consolidated into AlgorithmCollectionBuilder so all consumers get a default selector automatically.

How to Build and Test

# Build
mkdir -p build && cd build
cmake -DCMAKE_BUILD_TYPE=Release -DMSCCLPP_BUILD_EXT_TORCHCOMMS=ON ..
make -j$(nproc)

# Set env var
export TORCHCOMMS_BACKEND_LIB_PATH_MSCCLPP=$PWD/lib/_comms_mscclpp.cpython-*.so

# Run correctness tests
torchrun --nproc_per_node=8 test/torchcomms/test_correctness.py --all

# Run benchmarks
torchrun --nproc_per_node=8 test/torchcomms/bench_torchcomms.py --collective allreduce --warmup 100 --iters 200

[michael-beebe]

@microsoft-github-policy-service agree company="Microsoft"

[Copilot]

Pull request overview

Adds a TorchComms backend module (_comms_mscclpp*.so) that adapts TorchComms collective calls onto MSCCL++’s AlgorithmCollection (native + DSL), enabling PyTorch users to select MSCCL++ as a TorchComms backend at runtime.

Changes:

• Introduces a new C++ TorchComms backend implementation (bootstrap via c10d::Store, algorithm selection, executor + scratch, CUDA-event-based work tracking).
• Adds TorchComms-focused tests/benchmarks and a quickstart section documenting build/run steps.
• Adds a build option MSCCLPP_BUILD_EXT_TORCHCOMMS and a dedicated CMake target for building the backend module.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
CMakeLists.txt	Adds MSCCLPP_BUILD_EXT_TORCHCOMMS option and conditionally builds the TorchComms backend.
docs/quickstart.md	Documents how to build, use, test, and benchmark TorchComms support.
python/mscclpp_torchcomm/CMakeLists.txt	Fetches TorchComms sources/headers, builds _comms_mscclpp pybind module, links MSCCL++ + Torch.
python/mscclpp_torchcomm/__init__.py	Package stub for TorchComms backend directory.
python/mscclpp_torchcomm/requirements_cuda12.txt	Optional pip requirements for TorchComms backend environment.
python/mscclpp_torchcomm/csrc/TorchCommMSCCLPP.hpp	Declares the TorchComms backend class and supported/unsupported collectives.
python/mscclpp_torchcomm/csrc/TorchCommMSCCLPP.cpp	Implements init/finalize, algorithm selection wiring, and collective dispatch.
python/mscclpp_torchcomm/csrc/TorchCommMSCCLPPBootstrap.hpp	Declares bootstrap helper for rank/size + UniqueId exchange via store.
python/mscclpp_torchcomm/csrc/TorchCommMSCCLPPBootstrap.cpp	Implements UniqueId exchange and MSCCL++ communicator creation via TcpBootstrap.
python/mscclpp_torchcomm/csrc/TorchCommMSCCLPPPy.cpp	Exposes minimal pybind module + TorchComms dynamic loader entrypoint.
python/mscclpp_torchcomm/csrc/TorchWorkMSCCLPP.hpp	Declares CUDA event pool + TorchWork implementation.
python/mscclpp_torchcomm/csrc/TorchWorkMSCCLPP.cpp	Implements pooled CUDA events and async work completion tracking.
test/torchcomms/test_correctness.py	Correctness coverage for allreduce/allgather/reducescatter via TorchComms.
test/torchcomms/test_error_handling.py	Verifies clear runtime errors for unsupported ops and invalid usage patterns.
test/torchcomms/test_sizes.py	Sweeps message sizes to exercise selection boundaries and correctness.
test/torchcomms/test_training_loop.py	Simulates multi-iteration training-loop allreduce pattern.
test/torchcomms/test_user_algorithms.py	Validates user algorithm/selector registration via AlgorithmCollectionBuilder.
test/torchcomms/test_multicomm.py	Documents current multi-communicator limitation (expected skip).
test/torchcomms/bench_torchcomms.py	TorchComms benchmark driver for allreduce/allgather.
test/torchcomms/bench_report.py	Generates a report/figures from benchmark JSON output.
test/torchcomms/run_benchmarks.sh	Convenience runner to produce benchmark JSON + report + plots.