From 7f16ca9bab4df75f1ea2408531f6ea65c537aacd Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sat, 15 Aug 2026 19:07:25 +0800 Subject: [PATCH 1/3] feat(distributed): add deterministic TP8 all-reduce --- .../distributed/deterministic_collective.cu | 348 ++++++++++++++++++ csrc/ops.cpp | 34 ++ rl_engine/_C.pyi | 13 + rl_engine/distributed/__init__.py | 6 + rl_engine/distributed/collectives.py | 231 ++++++++++++ setup.py | 7 +- tests/distributed/__init__.py | 2 + .../test_deterministic_all_reduce.py | 97 +++++ 8 files changed, 737 insertions(+), 1 deletion(-) create mode 100644 csrc/cuda/distributed/deterministic_collective.cu create mode 100644 rl_engine/distributed/__init__.py create mode 100644 rl_engine/distributed/collectives.py create mode 100644 tests/distributed/__init__.py create mode 100644 tests/distributed/test_deterministic_all_reduce.py diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu new file mode 100644 index 00000000..d9b709d2 --- /dev/null +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -0,0 +1,348 @@ +// 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 kDeterministicWorldSize = 8; +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 4096; + +struct PeerPointers { + const void* values[kDeterministicWorldSize]; +}; + +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) { + const auto* rank0 = static_cast(peers.values[0]); + const auto* rank1 = static_cast(peers.values[1]); + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + 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 sum01 = ordered_add(rank0[index], rank1[index]); + const T sum23 = ordered_add(rank2[index], rank3[index]); + const T sum45 = ordered_add(rank4[index], rank5[index]); + const T sum67 = ordered_add(rank6[index], rank7[index]); + const T sum03 = ordered_add(sum01, sum23); + 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); + } +} + +class DeterministicCollectiveState { + public: + DeterministicCollectiveState( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) + : rank_(rank), + 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( + handles.size() == kDeterministicWorldSize, + "deterministic collectives require exactly 8 IPC handles"); + TORCH_CHECK( + offsets.size() == kDeterministicWorldSize, + "deterministic collectives require exactly 8 IPC offsets"); + TORCH_CHECK( + rank_ >= 0 && rank_ < kDeterministicWorldSize, + "deterministic collective rank must be in [0, 8)"); + + for (auto& peer : peers_.values) { + peer = nullptr; + } + imported_bases_.fill(nullptr); + try { + for (int peer = 0; peer < kDeterministicWorldSize; ++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: + deterministic_all_reduce_kernel<<>>( + peers_, + static_cast(output.data_ptr()), + element_count); + break; + case at::ScalarType::Half: + deterministic_all_reduce_kernel<<>>( + peers_, + static_cast(output.data_ptr()), + element_count); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + deterministic_all_reduce_kernel<<>>( + peers_, + static_cast(output.data_ptr()), + element_count); + 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 < kDeterministicWorldSize; ++peer) { + if (imported_bases_[peer] != nullptr) { + cudaIpcCloseMemHandle(imported_bases_[peer]); + imported_bases_[peer] = nullptr; + } + } + } + + int64_t rank_; + 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..cee0f19b 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -2,6 +2,19 @@ # 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..fbfd636b --- /dev/null +++ b/rl_engine/distributed/collectives.py @@ -0,0 +1,231 @@ +# 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 + +_TP8_WORLD_SIZE = 8 +_DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +_REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + + +class DeterministicCollective: + """Correctness-first deterministic CUDA collectives for one TP=8 node. + + The logical reduction tree is fixed to ``xor 1 -> xor 2 -> xor 4``: + ``((rank0 + rank1) + (rank2 + rank3)) + ((rank4 + rank5) + + (rank6 + rank7))``. Every node evaluates the lower logical rank before + the higher logical rank. + + 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 != _TP8_WORLD_SIZE: + raise ValueError( + f"deterministic collectives require world_size=8, 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 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. + """ + + 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..4ed51c34 --- /dev/null +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -0,0 +1,97 @@ +# 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 + +_WORLD_SIZE = 8 +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this TP=8 test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _WORLD_SIZE, + reason="requires eight visible CUDA GPUs", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + sum01 = values[0] + values[1] + sum23 = values[2] + values[3] + sum45 = values[4] + values[5] + sum67 = values[6] + values[7] + sum03 = sum01 + sum23 + sum47 = sum45 + sum67 + return sum03 + sum47 + + +def _all_gather_tensors(value: torch.Tensor) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(_WORLD_SIZE)] + dist.all_gather(gathered, value) + 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=_WORLD_SIZE, + timeout=timedelta(minutes=3), + ) + try: + with DeterministicCollective(device=device, max_size_bytes=1024 * 1024) as collective: + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator(device=device).manual_seed(20260815 + rank) + input = torch.randn(257, device=device, dtype=dtype, generator=generator) + peer_inputs = _all_gather_tensors(input) + expected = _fixed_tree_reference(peer_inputs) + + output = collective.all_reduce(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output) + 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) + 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_tp8_cuda() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_WORLD_SIZE, + join=True, + ) From 51f858500f1889f20c31b51cd157b333509850c6 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sat, 15 Aug 2026 21:04:07 +0800 Subject: [PATCH 2/3] feat(distributed): make all-reduce TP invariant --- .../distributed/deterministic_collective.cu | 136 +++++++++++++----- rl_engine/distributed/collectives.py | 23 +-- .../test_deterministic_all_reduce.py | 107 ++++++++------ 3 files changed, 179 insertions(+), 87 deletions(-) diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index d9b709d2..c7b9e6f4 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -21,14 +21,18 @@ namespace { -constexpr int kDeterministicWorldSize = 8; +constexpr int kMaxDeterministicWorldSize = 8; constexpr int kThreads = 256; constexpr int kMaxBlocks = 4096; struct PeerPointers { - const void* values[kDeterministicWorldSize]; + 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); @@ -53,27 +57,43 @@ __device__ __forceinline__ nv_bfloat16 ordered_add( } #endif -template -__device__ __forceinline__ T fixed_tree_reduce(const PeerPointers& peers, int64_t index) { +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]); - const auto* rank1 = static_cast(peers.values[1]); - const auto* rank2 = static_cast(peers.values[2]); - const auto* rank3 = static_cast(peers.values[3]); - 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 sum01 = ordered_add(rank0[index], rank1[index]); - const T sum23 = ordered_add(rank2[index], rank3[index]); - const T sum45 = ordered_add(rank4[index], rank5[index]); - const T sum67 = ordered_add(rank6[index], rank7[index]); - const T sum03 = ordered_add(sum01, sum23); - const T sum47 = ordered_add(sum45, sum67); - return ordered_add(sum03, sum47); + 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 +template __global__ void deterministic_all_reduce_kernel( PeerPointers peers, T* output, @@ -82,7 +102,37 @@ __global__ void deterministic_all_reduce_kernel( 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); + 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); } } @@ -94,6 +144,7 @@ class DeterministicCollectiveState { 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"); @@ -103,21 +154,24 @@ class DeterministicCollectiveState { "collective staging buffer must have dtype torch.uint8"); TORCH_CHECK(capacity_bytes_ > 0, "collective staging capacity must be positive"); TORCH_CHECK( - handles.size() == kDeterministicWorldSize, - "deterministic collectives require exactly 8 IPC handles"); + is_supported_world_size(world_size_), + "deterministic collectives require world size 1, 2, 4, or 8; got ", + world_size_); TORCH_CHECK( - offsets.size() == kDeterministicWorldSize, - "deterministic collectives require exactly 8 IPC offsets"); + offsets.size() == handles.size(), + "deterministic collectives require one IPC offset per handle"); TORCH_CHECK( - rank_ >= 0 && rank_ < kDeterministicWorldSize, - "deterministic collective rank must be in [0, 8)"); + 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 < kDeterministicWorldSize; ++peer) { + for (int peer = 0; peer < world_size_; ++peer) { TORCH_CHECK( handles[peer].size() == sizeof(cudaIpcMemHandle_t), "invalid CUDA IPC handle size for rank ", @@ -208,23 +262,32 @@ class DeterministicCollectiveState { switch (output.scalar_type()) { case at::ScalarType::Float: - deterministic_all_reduce_kernel<<>>( + launch_all_reduce( peers_, static_cast(output.data_ptr()), - element_count); + element_count, + blocks, + world_size_, + stream); break; case at::ScalarType::Half: - deterministic_all_reduce_kernel<<>>( + launch_all_reduce( peers_, static_cast(output.data_ptr()), - element_count); + element_count, + blocks, + world_size_, + stream); break; #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) case at::ScalarType::BFloat16: - deterministic_all_reduce_kernel<<>>( + launch_all_reduce( peers_, static_cast(output.data_ptr()), - element_count); + element_count, + blocks, + world_size_, + stream); break; #endif default: @@ -250,7 +313,7 @@ class DeterministicCollectiveState { } void close_imports() noexcept { - for (int peer = 0; peer < kDeterministicWorldSize; ++peer) { + for (int peer = 0; peer < world_size_; ++peer) { if (imported_bases_[peer] != nullptr) { cudaIpcCloseMemHandle(imported_bases_[peer]); imported_bases_[peer] = nullptr; @@ -259,13 +322,14 @@ class DeterministicCollectiveState { } 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_{}; + std::array imported_bases_{}; }; DeterministicCollectiveState* state_from_handle(int64_t handle) { diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index fbfd636b..4b000e32 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -11,18 +11,19 @@ import torch import torch.distributed as dist -_TP8_WORLD_SIZE = 8 +_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 deterministic CUDA collectives for one TP=8 node. + """Correctness-first TP-invariant CUDA collectives for one eight-GPU node. - The logical reduction tree is fixed to ``xor 1 -> xor 2 -> xor 4``: - ``((rank0 + rank1) + (rank2 + rank3)) + ((rank4 + rank5) + - (rank6 + rank7))``. Every node evaluates the lower logical rank before - the higher logical rank. + 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 @@ -47,9 +48,10 @@ def __init__( 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 != _TP8_WORLD_SIZE: + if self.world_size not in _SUPPORTED_WORLD_SIZES: raise ValueError( - f"deterministic collectives require world_size=8, got {self.world_size}" + "deterministic collectives require world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" ) if device is None: @@ -135,10 +137,11 @@ def all_reduce( *, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Return the fixed-tree sum on every rank. + """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. + ``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() diff --git a/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py index 4ed51c34..e0657202 100644 --- a/tests/distributed/test_deterministic_all_reduce.py +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -14,34 +14,36 @@ from rl_engine.distributed import DeterministicCollective -_WORLD_SIZE = 8 +_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 TP=8 test owns its worker processes; run pytest directly", + reason="this cross-TP test owns its worker processes; run pytest directly", ), pytest.mark.skipif( - torch.cuda.device_count() < _WORLD_SIZE, + torch.cuda.device_count() < _MAX_WORLD_SIZE, reason="requires eight visible CUDA GPUs", ), ] def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: - sum01 = values[0] + values[1] - sum23 = values[2] + values[3] - sum45 = values[4] + values[5] - sum67 = values[6] + values[7] - sum03 = sum01 + sum23 - sum47 = sum45 + sum67 - return sum03 + sum47 - - -def _all_gather_tensors(value: torch.Tensor) -> list[torch.Tensor]: - gathered = [torch.empty_like(value) for _ in range(_WORLD_SIZE)] - dist.all_gather(gathered, value) + 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 @@ -52,32 +54,55 @@ def _worker(rank: int, port: int) -> None: backend="nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, - world_size=_WORLD_SIZE, - timeout=timedelta(minutes=3), + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), ) try: - with DeterministicCollective(device=device, max_size_bytes=1024 * 1024) as collective: - for dtype in (torch.float32, torch.float16, torch.bfloat16): - generator = torch.Generator(device=device).manual_seed(20260815 + rank) - input = torch.randn(257, device=device, dtype=dtype, generator=generator) - peer_inputs = _all_gather_tensors(input) - expected = _fixed_tree_reference(peer_inputs) - - output = collective.all_reduce(input) - assert torch.equal(output, expected) - - peer_outputs = _all_gather_tensors(output) - 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) + 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() @@ -88,10 +113,10 @@ def _find_free_port() -> int: return int(sock.getsockname()[1]) -def test_deterministic_all_reduce_tp8_cuda() -> None: +def test_deterministic_all_reduce_cross_tp_cuda() -> None: mp.spawn( _worker, args=(_find_free_port(),), - nprocs=_WORLD_SIZE, + nprocs=_MAX_WORLD_SIZE, join=True, ) From 0b98cfd20d5dce3bfe95a83c9c6c871ea436a6af Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sat, 15 Aug 2026 21:16:03 +0800 Subject: [PATCH 3/3] style(distributed): format deterministic all-reduce --- rl_engine/_C.pyi | 1 - tests/distributed/test_deterministic_all_reduce.py | 12 +++--------- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index cee0f19b..a6ecaeec 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -14,7 +14,6 @@ def deterministic_collective_create( 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/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py index e0657202..eb406881 100644 --- a/tests/distributed/test_deterministic_all_reduce.py +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -58,9 +58,7 @@ def _worker(rank: int, port: int) -> None: timeout=timedelta(minutes=5), ) try: - groups = { - tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES - } + 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( @@ -80,18 +78,14 @@ def _worker(rank: int, port: int) -> None: generator=generator, ).to(device=device, dtype=dtype) leaves = list(leaves_tensor.unbind()) - input = _fixed_tree_reference( - leaves[start : start + leaves_per_rank] - ) + 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 - ) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) baseline = output.clone() for _ in range(3):