diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu new file mode 100644 index 00000000..c7b9e6f4 --- /dev/null +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -0,0 +1,412 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kMaxDeterministicWorldSize = 8; +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 4096; + +struct PeerPointers { + const void* values[kMaxDeterministicWorldSize]; +}; + +bool is_supported_world_size(int64_t world_size) { + return world_size == 1 || world_size == 2 || world_size == 4 || world_size == 8; +} + +template +__device__ __forceinline__ T ordered_add(T lower, T upper); + +template <> +__device__ __forceinline__ float ordered_add(float lower, float upper) { + float result; + asm volatile("add.rn.f32 %0, %1, %2;" : "=f"(result) : "f"(lower), "f"(upper)); + return result; +} + +template <> +__device__ __forceinline__ half ordered_add(half lower, half upper) { + return __hadd(lower, upper); +} + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template <> +__device__ __forceinline__ nv_bfloat16 ordered_add( + nv_bfloat16 lower, + nv_bfloat16 upper) { + return __hadd(lower, upper); +} +#endif + +template +__device__ __forceinline__ T fixed_tree_reduce( + const PeerPointers& peers, + int64_t index) { + static_assert( + WorldSize == 1 || WorldSize == 2 || WorldSize == 4 || WorldSize == 8, + "deterministic collectives only support TP sizes 1, 2, 4, and 8"); + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const T sum01 = ordered_add(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const T sum23 = ordered_add(rank2[index], rank3[index]); + const T sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const T sum45 = ordered_add(rank4[index], rank5[index]); + const T sum67 = ordered_add(rank6[index], rank7[index]); + const T sum47 = ordered_add(sum45, sum67); + return ordered_add(sum03, sum47); + } + } + } +} + +template +__global__ void deterministic_all_reduce_kernel( + PeerPointers peers, + T* output, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = fixed_tree_reduce(peers, index); + } +} + +template +void launch_all_reduce( + const PeerPointers& peers, + T* output, + int64_t element_count, + int blocks, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 2: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 4: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 8: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + +class DeterministicCollectiveState { + public: + DeterministicCollectiveState( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) + : rank_(rank), + world_size_(handles.size()), + device_index_(staging.get_device()), + capacity_bytes_(staging.numel() * staging.element_size()) { + TORCH_CHECK(staging.is_cuda(), "collective staging buffer must be CUDA"); + TORCH_CHECK(staging.is_contiguous(), "collective staging buffer must be contiguous"); + TORCH_CHECK( + staging.scalar_type() == torch::kUInt8, + "collective staging buffer must have dtype torch.uint8"); + TORCH_CHECK(capacity_bytes_ > 0, "collective staging capacity must be positive"); + TORCH_CHECK( + is_supported_world_size(world_size_), + "deterministic collectives require world size 1, 2, 4, or 8; got ", + world_size_); + TORCH_CHECK( + offsets.size() == handles.size(), + "deterministic collectives require one IPC offset per handle"); + TORCH_CHECK( + rank_ >= 0 && rank_ < world_size_, + "deterministic collective rank must be in [0, ", + world_size_, + ")"); + + for (auto& peer : peers_.values) { + peer = nullptr; + } + imported_bases_.fill(nullptr); + try { + for (int peer = 0; peer < world_size_; ++peer) { + TORCH_CHECK( + handles[peer].size() == sizeof(cudaIpcMemHandle_t), + "invalid CUDA IPC handle size for rank ", + peer); + TORCH_CHECK(offsets[peer] >= 0, "negative CUDA IPC offset for rank ", peer); + + if (peer == rank_) { + peers_.values[peer] = staging.data_ptr(); + continue; + } + + cudaIpcMemHandle_t handle{}; + auto* raw_handle = reinterpret_cast(&handle); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + TORCH_CHECK( + handles[peer][byte] >= 0 && handles[peer][byte] <= 255, + "invalid CUDA IPC handle byte for rank ", + peer); + raw_handle[byte] = static_cast(handles[peer][byte]); + } + + void* base = nullptr; + AT_CUDA_CHECK(cudaIpcOpenMemHandle( + &base, + handle, + cudaIpcMemLazyEnablePeerAccess)); + imported_bases_[peer] = base; + peers_.values[peer] = static_cast(base) + offsets[peer]; + } + } catch (...) { + close_imports(); + throw; + } + } + + ~DeterministicCollectiveState() { + int previous_device = -1; + if (cudaGetDevice(&previous_device) == cudaSuccess && previous_device != device_index_) { + if (cudaSetDevice(device_index_) != cudaSuccess) { + return; + } + } + close_imports(); + if (previous_device >= 0 && previous_device != device_index_) { + cudaSetDevice(previous_device); + } + } + + void stage(torch::Tensor& input, cudaStream_t stream) { + check_tensor(input, "input"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK( + input_bytes <= capacity_bytes_, + "input requires ", + input_bytes, + " bytes but staging capacity is ", + capacity_bytes_); + if (input_bytes > 0) { + AT_CUDA_CHECK(cudaMemcpyAsync( + const_cast(peers_.values[rank_]), + input.data_ptr(), + input_bytes, + cudaMemcpyDeviceToDevice, + stream)); + } + staged_bytes_ = input_bytes; + staged_scalar_type_ = input.scalar_type(); + has_staged_input_ = true; + } + + void all_reduce(torch::Tensor& output, cudaStream_t stream) const { + check_tensor(output, "output"); + TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); + TORCH_CHECK( + output.scalar_type() == staged_scalar_type_, + "all-reduce output dtype must match the staged input dtype"); + TORCH_CHECK( + output.numel() * output.element_size() == staged_bytes_, + "all-reduce output size must match the staged input size"); + + const int64_t element_count = output.numel(); + if (element_count == 0) { + return; + } + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic all-reduce supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } + + private: + void check_tensor(const torch::Tensor& tensor, const char* name) const { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK( + tensor.get_device() == device_index_, + name, + " must be on cuda:", + device_index_, + ", got ", + tensor.device()); + } + + void close_imports() noexcept { + for (int peer = 0; peer < world_size_; ++peer) { + if (imported_bases_[peer] != nullptr) { + cudaIpcCloseMemHandle(imported_bases_[peer]); + imported_bases_[peer] = nullptr; + } + } + } + + int64_t rank_; + int64_t world_size_; + int device_index_; + int64_t capacity_bytes_; + int64_t staged_bytes_{0}; + at::ScalarType staged_scalar_type_{at::ScalarType::Undefined}; + bool has_staged_input_{false}; + PeerPointers peers_{}; + std::array imported_bases_{}; +}; + +DeterministicCollectiveState* state_from_handle(int64_t handle) { + TORCH_CHECK(handle != 0, "deterministic collective handle is closed"); + return reinterpret_cast(handle); +} + +} // namespace + +std::tuple, int64_t> deterministic_collective_ipc_meta( + torch::Tensor& tensor) { + const c10::cuda::CUDAGuard device_guard(tensor.device()); + TORCH_CHECK(tensor.is_cuda(), "IPC tensor must be CUDA"); + TORCH_CHECK(tensor.is_contiguous(), "IPC tensor must be contiguous"); + TORCH_CHECK(tensor.numel() > 0, "cannot export an empty CUDA allocation"); + + CUdeviceptr allocation_base = 0; + size_t allocation_size = 0; + const auto pointer = reinterpret_cast(tensor.data_ptr()); + TORCH_CHECK( + cuPointerGetAttribute( + &allocation_base, + CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, + pointer) == CUDA_SUCCESS, + "failed to query CUDA allocation base"); + TORCH_CHECK( + cuPointerGetAttribute( + &allocation_size, + CU_POINTER_ATTRIBUTE_RANGE_SIZE, + pointer) == CUDA_SUCCESS, + "failed to query CUDA allocation size"); + + const int64_t offset = static_cast(pointer - allocation_base); + const int64_t tensor_bytes = tensor.numel() * tensor.element_size(); + TORCH_CHECK(offset >= 0, "invalid negative CUDA allocation offset"); + TORCH_CHECK( + static_cast(offset + tensor_bytes) <= allocation_size, + "IPC tensor exceeds its CUDA allocation"); + + cudaIpcMemHandle_t handle{}; + AT_CUDA_CHECK(cudaIpcGetMemHandle( + &handle, + reinterpret_cast(allocation_base))); + const auto* raw_handle = reinterpret_cast(&handle); + std::vector bytes(sizeof(handle)); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + bytes[byte] = raw_handle[byte]; + } + return std::make_tuple(bytes, offset); +} + +int64_t deterministic_collective_create( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) { + const c10::cuda::CUDAGuard device_guard(staging.device()); + auto state = std::make_unique( + staging, + handles, + offsets, + rank); + return reinterpret_cast(state.release()); +} + +void deterministic_collective_destroy(int64_t handle) { + delete state_from_handle(handle); +} + +void deterministic_collective_stage(int64_t handle, torch::Tensor& input) { + const c10::cuda::CUDAGuard device_guard(input.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->stage(input, stream); +} + +void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->all_reduce(output, stream); +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index eee328a4..87362535 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -88,6 +88,18 @@ torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tenso torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); +// Single-node TP=8 deterministic collectives. +std::tuple, int64_t> deterministic_collective_ipc_meta( + torch::Tensor& tensor); +int64_t deterministic_collective_create( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank); +void deterministic_collective_destroy(int64_t handle); +void deterministic_collective_stage(int64_t handle, torch::Tensor& input); +void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output); + // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); @@ -355,6 +367,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); + // Single-node TP=8 fixed-tree collectives. + m.def( + "deterministic_collective_ipc_meta", + &deterministic_collective_ipc_meta, + "Export a CUDA allocation for deterministic collectives"); + m.def( + "deterministic_collective_create", + &deterministic_collective_create, + "Create a single-node TP=8 deterministic collective state"); + m.def( + "deterministic_collective_destroy", + &deterministic_collective_destroy, + "Destroy a deterministic collective state"); + m.def( + "deterministic_collective_stage", + &deterministic_collective_stage, + "Stage one rank's input in symmetric CUDA IPC memory"); + m.def( + "deterministic_collective_all_reduce", + &deterministic_collective_all_reduce, + "Run the TP=8 deterministic fixed-tree all-reduce kernel"); + // registry Prefix-Shared Attention m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 20f17461..a6ecaeec 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -2,6 +2,18 @@ # This file is a type stub for the compiled C++ extension module. import torch +def deterministic_collective_ipc_meta( + tensor: torch.Tensor, +) -> tuple[list[int], int]: ... +def deterministic_collective_create( + staging: torch.Tensor, + handles: list[list[int]], + offsets: list[int], + rank: int, +) -> int: ... +def deterministic_collective_destroy(handle: int) -> None: ... +def deterministic_collective_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_all_reduce(handle: int, output: torch.Tensor) -> None: ... def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... def batch_invariant_logp_sm90( diff --git a/rl_engine/distributed/__init__.py b/rl_engine/distributed/__init__.py new file mode 100644 index 00000000..37698f1a --- /dev/null +++ b/rl_engine/distributed/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from rl_engine.distributed.collectives import DeterministicCollective + +__all__ = ["DeterministicCollective"] diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py new file mode 100644 index 00000000..4b000e32 --- /dev/null +++ b/rl_engine/distributed/collectives.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import socket +import threading +from types import TracebackType +from typing import Any + +import torch +import torch.distributed as dist + +_SUPPORTED_WORLD_SIZES = (1, 2, 4, 8) +_DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +_REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + + +class DeterministicCollective: + """Correctness-first TP-invariant CUDA collectives for one eight-GPU node. + + TP sizes 1, 2, 4, and 8 use nested prefixes of the same balanced tree. + A reduction is cross-TP bitwise invariant when every rank input is the + corresponding contiguous subtree root of one canonical finest-grained + reduction, as produced by a TBIK-compatible row-parallel kernel. Every + node evaluates the lower logical subtree before the higher one. + + One instance owns a symmetric CUDA IPC staging buffer. All ranks must call + its methods in the same order with matching shapes and dtypes. Calls are + host-synchronizing by design; the first version prioritizes determinism and + lifetime safety over overlap or throughput. + """ + + def __init__( + self, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, + ) -> None: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before collectives") + if not torch.cuda.is_available(): + raise RuntimeError("deterministic collectives require CUDA") + if max_size_bytes <= 0: + raise ValueError("max_size_bytes must be positive") + + self.group = group if group is not None else dist.group.WORLD + self.rank = dist.get_rank(group=self.group) + self.world_size = dist.get_world_size(group=self.group) + if self.world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "deterministic collectives require world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" + ) + + if device is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + elif isinstance(device, int): + normalized_device = torch.device("cuda", device) + else: + normalized_device = torch.device(device) + if normalized_device.type != "cuda": + raise ValueError(f"deterministic collectives require a CUDA device, got {device!r}") + if normalized_device.index is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + if normalized_device.index != torch.cuda.current_device(): + raise ValueError( + "the collective device must be the current CUDA device; call " + f"torch.cuda.set_device({normalized_device.index}) first" + ) + + try: + from rl_engine import _C + except ImportError as exc: + raise RuntimeError( + "the RL-Kernel CUDA extension is required; rebuild with " + "`pip install --no-build-isolation -e .`" + ) from exc + required_symbols = ( + "deterministic_collective_ipc_meta", + "deterministic_collective_create", + "deterministic_collective_destroy", + "deterministic_collective_stage", + "deterministic_collective_all_reduce", + ) + missing = [name for name in required_symbols if not hasattr(_C, name)] + if missing: + raise RuntimeError( + "the RL-Kernel CUDA extension lacks deterministic collectives: " + + ", ".join(missing) + ) + + self.device = normalized_device + self.max_size_bytes = int(max_size_bytes) + self._extension = _C + self._lock = threading.Lock() + self._handle = 0 + self._staging = torch.empty( + self.max_size_bytes, + dtype=torch.uint8, + device=self.device, + ) + + handle, offset = self._extension.deterministic_collective_ipc_meta(self._staging) + local_meta = { + "handle": handle, + "offset": int(offset), + "capacity": self.max_size_bytes, + "hostname": socket.gethostname(), + } + gathered_meta: list[dict[str, Any] | None] = [None] * self.world_size + dist.all_gather_object(gathered_meta, local_meta, group=self.group) + if any(meta is None for meta in gathered_meta): + raise RuntimeError("failed to exchange CUDA IPC metadata") + complete_meta = [meta for meta in gathered_meta if meta is not None] + hostnames = {meta["hostname"] for meta in complete_meta} + if len(hostnames) != 1: + raise ValueError("deterministic collectives require all ranks on one host") + capacities = {meta["capacity"] for meta in complete_meta} + if capacities != {self.max_size_bytes}: + raise ValueError("all ranks must use the same max_size_bytes") + + handles = [meta["handle"] for meta in complete_meta] + offsets = [meta["offset"] for meta in complete_meta] + self._handle = self._extension.deterministic_collective_create( + self._staging, + handles, + offsets, + self.rank, + ) + self._synchronize_ranks() + + def all_reduce( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the TBIK-compatible fixed-tree sum on every rank. + + Supported dtypes are float32, float16, and bfloat16. ``out`` may alias + ``input``; the input is staged before the output kernel starts. Cross-TP + invariance requires inputs to follow the class-level subtree contract. + """ + + self._check_open() + self._validate_reduction_input(input) + if out is None: + out = torch.empty_like(input) + self._validate_output(out, input) + + with self._lock: + self._validate_matching_signature("all_reduce", input) + self._extension.deterministic_collective_stage(self._handle, input) + self._synchronize_ranks() + self._extension.deterministic_collective_all_reduce(self._handle, out) + self._synchronize_ranks() + return out + + def close(self) -> None: + """Release imported CUDA IPC mappings after the last collective call.""" + + handle = getattr(self, "_handle", 0) + if not handle: + return + torch.cuda.synchronize(self.device) + self._handle = 0 + self._extension.deterministic_collective_destroy(handle) + + def __enter__(self) -> DeterministicCollective: + self._check_open() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def _check_open(self) -> None: + if not getattr(self, "_handle", 0): + raise RuntimeError("deterministic collective is closed") + + def _validate_reduction_input(self, input: torch.Tensor) -> None: + if not input.is_cuda or input.device != self.device: + raise ValueError(f"input must be on {self.device}, got {input.device}") + if not input.is_contiguous(): + raise ValueError("input must be contiguous") + if input.dtype not in _REDUCTION_DTYPES: + raise TypeError( + "deterministic reductions support float32, float16, and bfloat16; " + f"got {input.dtype}" + ) + input_bytes = input.numel() * input.element_size() + if input_bytes > self.max_size_bytes: + raise ValueError( + f"input requires {input_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + + def _validate_output(self, output: torch.Tensor, input: torch.Tensor) -> None: + if output.device != input.device: + raise ValueError("out must be on the same device as input") + if output.dtype != input.dtype: + raise TypeError("out must have the same dtype as input") + if output.shape != input.shape: + raise ValueError("out must have the same shape as input") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: + signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + signatures: list[tuple[Any, ...] | None] = [None] * self.world_size + dist.all_gather_object(signatures, signature, group=self.group) + if any(peer_signature != signature for peer_signature in signatures): + raise ValueError( + f"all ranks must call {op_name} with matching shapes and dtypes; got {signatures}" + ) + + def _synchronize_ranks(self) -> None: + torch.cuda.synchronize(self.device) + backend = dist.get_backend(self.group) + if backend == dist.Backend.NCCL or str(backend).lower() == "nccl": + dist.barrier(group=self.group, device_ids=[self.device.index]) + else: + dist.barrier(group=self.group) diff --git a/setup.py b/setup.py index a17ecb40..2c7af89e 100644 --- a/setup.py +++ b/setup.py @@ -83,6 +83,7 @@ def get_extensions(): "csrc/cuda/rmsnorm.cu", "csrc/cuda/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", ] cc_major, cc_minor = torch.cuda.get_device_capability() @@ -147,6 +148,9 @@ def get_extensions(): cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") sm90_srcs = [ "csrc/cuda/fused_logp_sm90.cu", @@ -162,7 +166,8 @@ def get_extensions(): cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - extra_link_args.append("-lcuda") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in diff --git a/tests/distributed/__init__.py b/tests/distributed/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/tests/distributed/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py new file mode 100644 index 00000000..eb406881 --- /dev/null +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _MAX_WORLD_SIZE, + reason="requires eight visible CUDA GPUs", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + level = values + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _worker(rank: int, port: int) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + leaves_per_rank = _MAX_WORLD_SIZE // tp_size + start = group_rank * leaves_per_rank + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator().manual_seed(20260815) + leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + 257, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + leaves = list(leaves_tensor.unbind()) + input = _fixed_tree_reference(leaves[start : start + leaves_per_rank]) + expected = _fixed_tree_reference(leaves) + + output = collective.all_reduce(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) + + baseline = output.clone() + for _ in range(3): + repeated = collective.all_reduce(input) + assert torch.equal(repeated, baseline) + + inplace = input.clone() + returned = collective.all_reduce(inplace, out=inplace) + assert returned is inplace + assert torch.equal(inplace, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_all_reduce_cross_tp_cuda() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + )