From 5cf23ecc5882537edddf1c7d7282f4c64dddf521 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 16:42:19 +0000 Subject: [PATCH 1/5] feat(distributed): add deterministic ROCm collectives --- benchmarks/benchmark_rocm_collectives.py | 290 ++++++++++++ csrc/ops.cpp | 14 +- docs/design/rocm-deterministic-collectives.md | 112 +++++ rl_engine/distributed/__init__.py | 12 +- .../distributed/transport_collectives.py | 419 ++++++++++++++++++ .../kernels/ops/cuda/attention/__init__.py | 2 + .../kernels/ops/cuda/attention/cp_comm.py | 183 +++++++- rl_engine/kernels/ops/pytorch/ffn/ffn.py | 39 +- setup.py | 9 +- .../test_rocm_attention_transport.py | 95 ++++ ...test_transport_deterministic_collective.py | 375 ++++++++++++++++ tests/test_build_platform_collectives.py | 54 +++ tests/test_qwen_ffn.py | 19 + tests/test_rocm_collective_benchmark.py | 73 +++ 14 files changed, 1676 insertions(+), 20 deletions(-) create mode 100644 benchmarks/benchmark_rocm_collectives.py create mode 100644 docs/design/rocm-deterministic-collectives.md create mode 100644 rl_engine/distributed/transport_collectives.py create mode 100644 tests/distributed/test_rocm_attention_transport.py create mode 100644 tests/distributed/test_transport_deterministic_collective.py create mode 100644 tests/test_build_platform_collectives.py create mode 100644 tests/test_rocm_collective_benchmark.py diff --git a/benchmarks/benchmark_rocm_collectives.py b/benchmarks/benchmark_rocm_collectives.py new file mode 100644 index 00000000..6202e2e5 --- /dev/null +++ b/benchmarks/benchmark_rocm_collectives.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark deterministic ROCm collectives against native RCCL. + +Example: + + torchrun --standalone --nproc-per-node=8 \ + benchmarks/benchmark_rocm_collectives.py \ + --size-bytes 4096 65536 1048576 16777216 \ + --output benchmarks/results/rocm_collectives_mi300x.json + +The native RCCL rows are performance references only. They are not used as a +bitwise correctness oracle because their floating-point reduction order is not +part of the strict deterministic contract. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +from pathlib import Path +from typing import Callable, Sequence + +import torch +import torch.distributed as dist + +from rl_engine.distributed import RCCLDeterministicCollective + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +_OPERATIONS = ("all_reduce", "all_gather", "reduce_scatter") + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--size-bytes", + type=int, + nargs="+", + default=[4 * 1024, 64 * 1024, 1024 * 1024, 16 * 1024 * 1024], + ) + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") + parser.add_argument("--operations", nargs="+", choices=_OPERATIONS, default=_OPERATIONS) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--samples", type=int, default=5) + parser.add_argument("--output", type=Path) + return parser.parse_args(argv) + + +def _validate_args(args: argparse.Namespace) -> None: + if any(size <= 0 for size in args.size_bytes): + raise ValueError("every --size-bytes value must be positive") + if args.warmup < 0: + raise ValueError("--warmup must be non-negative") + if args.iterations <= 0 or args.samples <= 0: + raise ValueError("--iterations and --samples must be positive") + + +def _timed_sample(operation: Callable[[], None], *, warmup: int, iterations: int) -> float: + for _ in range(warmup): + operation() + torch.cuda.synchronize() + dist.barrier() + start = time.perf_counter() + for _ in range(iterations): + operation() + torch.cuda.synchronize() + elapsed = (time.perf_counter() - start) / iterations + + # Report the slowest rank, which is the end-to-end collective latency. + elapsed_tensor = torch.tensor([elapsed], dtype=torch.float64, device="cuda") + dist.all_reduce(elapsed_tensor, op=dist.ReduceOp.MAX) + return float(elapsed_tensor.item()) + + +def _benchmark( + operation: Callable[[], None], + *, + warmup: int, + iterations: int, + samples: int, +) -> dict[str, object]: + timings = [ + _timed_sample(operation, warmup=warmup if index == 0 else 0, iterations=iterations) + for index in range(samples) + ] + median = statistics.median(timings) + return { + "median_us": median * 1.0e6, + "min_us": min(timings) * 1.0e6, + "max_us": max(timings) * 1.0e6, + "samples_us": [value * 1.0e6 for value in timings], + } + + +def _make_inputs( + *, + size_bytes: int, + dtype: torch.dtype, + world_size: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, int]: + element_size = torch.empty((), dtype=dtype).element_size() + elements = max(world_size, size_bytes // element_size) + elements -= elements % world_size + generator = torch.Generator(device="cpu").manual_seed(942 + rank) + tensor = torch.randn(elements, generator=generator, dtype=torch.float32).to( + device=device, + dtype=dtype, + ) + return tensor.contiguous(), elements * element_size + + +def _operation_pair( + name: str, + input_tensor: torch.Tensor, + collective: RCCLDeterministicCollective, + world_size: int, +) -> tuple[Callable[[], None], Callable[[], None], torch.Tensor, torch.Tensor]: + if name == "all_reduce": + deterministic_out = torch.empty_like(input_tensor) + native_out = torch.empty_like(input_tensor) + + def deterministic() -> None: + collective.all_reduce(input_tensor, out=deterministic_out) + + def native() -> None: + native_out.copy_(input_tensor) + dist.all_reduce(native_out) + + elif name == "all_gather": + output_shape = (input_tensor.numel() * world_size,) + deterministic_out = torch.empty(output_shape, dtype=input_tensor.dtype, device="cuda") + native_out = torch.empty_like(deterministic_out) + + def deterministic() -> None: + collective.all_gather(input_tensor, out=deterministic_out) + + def native() -> None: + dist.all_gather_into_tensor(native_out, input_tensor) + + elif name == "reduce_scatter": + output_shape = (input_tensor.numel() // world_size,) + deterministic_out = torch.empty(output_shape, dtype=input_tensor.dtype, device="cuda") + native_out = torch.empty_like(deterministic_out) + + def deterministic() -> None: + collective.reduce_scatter(input_tensor, out=deterministic_out) + + def native() -> None: + dist.reduce_scatter_tensor(native_out, input_tensor) + + else: # pragma: no cover - argparse constrains this value + raise ValueError(f"unsupported operation: {name}") + + return deterministic, native, deterministic_out, native_out + + +def run(args: argparse.Namespace) -> dict[str, object] | None: + _validate_args(args) + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("the ROCm collective benchmark requires an available AMD GPU") + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", init_method="env://") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size not in (2, 4, 8): + raise RuntimeError(f"the benchmark requires 2, 4, or 8 ranks, got {world_size}") + device = torch.device("cuda", local_rank) + dtype = _DTYPES[args.dtype] + max_size_bytes = max(args.size_bytes) + dtype.itemsize * world_size + + rows: list[dict[str, object]] = [] + try: + with RCCLDeterministicCollective( + device=device, + max_size_bytes=max_size_bytes, + ) as collective: + for requested_size in args.size_bytes: + input_tensor, actual_size = _make_inputs( + size_bytes=requested_size, + dtype=dtype, + world_size=world_size, + rank=rank, + device=device, + ) + for name in args.operations: + deterministic, native, deterministic_out, native_out = _operation_pair( + name, + input_tensor, + collective, + world_size, + ) + deterministic() + deterministic_repeat = deterministic_out.clone() + deterministic() + repeat_bitwise = bool(torch.equal(deterministic_out, deterministic_repeat)) + native() + max_abs_vs_native = float( + (deterministic_out.float() - native_out.float()).abs().max().item() + ) + + torch.cuda.reset_peak_memory_stats(device) + deterministic_timing = _benchmark( + deterministic, + warmup=args.warmup, + iterations=args.iterations, + samples=args.samples, + ) + deterministic_peak = int(torch.cuda.max_memory_allocated(device)) + torch.cuda.reset_peak_memory_stats(device) + native_timing = _benchmark( + native, + warmup=args.warmup, + iterations=args.iterations, + samples=args.samples, + ) + native_peak = int(torch.cuda.max_memory_allocated(device)) + deterministic_us = float(deterministic_timing["median_us"]) + native_us = float(native_timing["median_us"]) + rows.append( + { + "operation": name, + "requested_size_bytes": requested_size, + "actual_input_bytes": actual_size, + "dtype": args.dtype, + "deterministic": deterministic_timing, + "native_rccl": native_timing, + "latency_ratio_vs_native": deterministic_us / native_us, + "deterministic_input_gbps": actual_size / (deterministic_us * 1.0e3), + "native_input_gbps": actual_size / (native_us * 1.0e3), + "repeat_bitwise": repeat_bitwise, + "max_abs_vs_native": max_abs_vs_native, + "deterministic_workspace_bytes": collective.workspace_size_bytes, + "deterministic_peak_allocated_bytes": deterministic_peak, + "native_peak_allocated_bytes": native_peak, + } + ) + + reports: list[list[dict[str, object]] | None] = [None] * world_size + dist.all_gather_object(reports, rows) + if rank != 0: + return None + payload = { + "schema_version": "rlkernel.rocm_collective_benchmark.v1", + "world_size": world_size, + "device": torch.cuda.get_device_name(device), + "hip_version": torch.version.hip, + "collective_backend": RCCLDeterministicCollective.backend_id, + "reduction_order": RCCLDeterministicCollective.reduction_order, + "supports_compute_communication_fusion": False, + "warmup": args.warmup, + "iterations": args.iterations, + "samples": args.samples, + "rows": rows, + "all_rank_repeat_bitwise": all( + bool(row["repeat_bitwise"]) + for rank_rows in reports + if rank_rows is not None + for row in rank_rows + ), + } + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8") + return payload + finally: + dist.destroy_process_group() + + +def main(argv: Sequence[str] | None = None) -> int: + payload = run(parse_args(argv)) + if payload is not None: + print(json.dumps(payload, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f1e66dba..fb84dc64 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -78,7 +78,7 @@ torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::optional bias); #endif -#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_ROCM) torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); @@ -93,7 +93,9 @@ 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. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) +// Single-node TP=8 deterministic CUDA IPC collectives. ROCm uses the +// rank-ordered RCCL transport in rl_engine.distributed.transport_collectives. std::tuple, int64_t> deterministic_collective_ipc_meta( torch::Tensor& tensor); int64_t deterministic_collective_create( @@ -106,6 +108,7 @@ void deterministic_collective_stage(int64_t handle, torch::Tensor& input); void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output); void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& output); void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); +#endif // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); @@ -391,7 +394,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Single-card SM90 batch-invariant LM-head forward with fp32 output"); #endif -#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_ROCM) m.def("fused_logp_forward_out", &fused_logp_forward_out, "Fused logp out"); m.def("fused_logp_forward_fp32", &fused_logp_forward_fp32, "Fused logp fp32"); m.def("fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out, "Fused logp indexed out"); @@ -406,7 +409,9 @@ 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. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) + // Single-node TP=8 fixed-tree CUDA IPC collectives. ROCm dispatches to + // the Python RCCL transport implementation instead. m.def( "deterministic_collective_ipc_meta", &deterministic_collective_ipc_meta, @@ -435,6 +440,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_collective_all_gather", &deterministic_collective_all_gather, "Run the TP=8 deterministic rank-ordered all-gather kernel"); +#endif // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. #if !defined(USE_ROCM) diff --git a/docs/design/rocm-deterministic-collectives.md b/docs/design/rocm-deterministic-collectives.md new file mode 100644 index 00000000..ce104332 --- /dev/null +++ b/docs/design/rocm-deterministic-collectives.md @@ -0,0 +1,112 @@ +# ROCm deterministic collectives + +This document defines the ROCm communication boundary used by deterministic +TP/CP kernels. The implementation is intentionally separate from the CUDA IPC +collective in `csrc/cuda/distributed/deterministic_collective.cu`. + +## Contract + +`RCCLDeterministicCollective` implements `all_reduce`, `all_gather`, and +`reduce_scatter` with the same Python call shape as the CUDA collective. It +supports process-group sizes 1, 2, 4, and 8 and FP32/FP16/BF16 reductions. +Inputs and outputs must be contiguous and all ranks must call operations in the +same order with matching shapes, dtypes, and capacity. + +RCCL is used only for rank-ordered tensor transport: + +1. `all_gather_into_tensor` gathers every rank's bit patterns. +2. Each rank evaluates the same balanced tree locally: + `((rank0 + rank1) + (rank2 + rank3)) + ...`. +3. `reduce_scatter` slices the rank-owned rows after that fixed reduction. + +RCCL's `all_reduce` and `reduce_scatter` are not used for strict reductions. +They guarantee a mathematical reduction but do not expose a stable +floating-point operand order. Delegating the arithmetic to them would weaken +the cross-TP bitwise contract. + +`create_deterministic_collective` is the platform boundary. It selects the +existing CUDA IPC implementation on NVIDIA and the RCCL transport +implementation when `torch.version.hip` is set. The FFN path calls this factory +instead of importing the CUDA class directly. + +The ROCm extension build also excludes the CUDA IPC source and does not link +`libcuda`; the Python transport has no CUDA-driver dependency. + +## Relationship to vLLM + +The backend split follows the useful parts of vLLM's device communicator +design: + +- PyTorch exposes RCCL through the `nccl` process-group API. +- ROCm AllGather uses `torch.distributed.all_gather_into_tensor` rather than a + manually allocated PyNccl path. +- Backend, topology, world-size, dtype, layout, and capacity checks fail closed + before entering an optimized path. + +vLLM QuickReduce, custom all-reduce, and AITER all-reduce are not used in the +strict path. Those are valuable performance implementations, but their +reduction order is not the fixed balanced tree required here. They can be +added later as an explicitly non-strict performance mode with separate +provenance and toleranced correctness tests. + +## Compute/communication fusion + +The current ROCm collective is synchronous and reports +`supports_async_overlap = False` and +`supports_compute_communication_fusion = False`. FFN and Attention keep the +dependency boundaries explicit: + +```text +sequence-parallel FFN: AllGather(input) -> GEMM -> ReduceScatter(output) +strict CP Attention: AllGather(Q/K/V/positions) -> Attention -> Scatter(output/LSE) +``` + +This is neither a fused GEMM+collective kernel nor two-stream overlap. vLLM's +generic AsyncTP GEMM/communication fusion is currently a CUDA path; ROCm AITER +has narrower fusions such as all-reduce plus RMSNorm, which do not replace this +contract. + +Future overlap must preserve collective issue order, the local reduction tree, +and stage boundaries. It also needs repeat-bitwise tests before being marked +strict. A useful first candidate is backward work that is independent of a +pending reduction; the forward SP AllGather and final row-parallel reduction +are data dependencies and cannot simply be overlapped with their adjacent +GEMMs. + +The “fill rank slots as messages arrive and merge ready contiguous blocks” +scheme needs a lower-level P2P or HIP/XGMI transport. PyTorch/RCCL +`all_gather_into_tensor` exposes completion of the whole collective, not +per-rank arrival events. Such a pipeline may merge only canonical sibling +subtrees when both are ready; merging arbitrary contiguous arrivals would +change floating-point parenthesization. It is therefore a follow-up transport, +not an optimization silently hidden inside this baseline. + +## Performance acceptance + +The transport-only baseline favors correctness and portability. AllGather +writes directly to its final output; reductions reuse one lazily grown +`world_size * input_bytes` byte workspace until `close()`. It still moves more +data than a native RCCL AllReduce. A ROCm GPU PR should therefore report, for +world sizes 2/4/8 and representative FFN tensors: + +- latency and effective bandwidth for all three collectives; +- peak temporary memory; +- comparison with RCCL and vLLM's available ROCm communicator; +- repeat-bitwise and cross-TP results; +- end-to-end TP/CP/SP FFN timing, not only isolated transport timing. + +A later fixed-tree HIP/XGMI implementation may replace the transport behind the +same factory after it satisfies those checks. + +Run the included native-RCCL comparison on a single node, for example: + +```bash +torchrun --standalone --nproc-per-node=8 \ + benchmarks/benchmark_rocm_collectives.py \ + --size-bytes 4096 65536 1048576 16777216 \ + --output benchmarks/results/rocm_collectives_mi300x.json +``` + +The benchmark records slowest-rank latency, temporary allocation, repeat +bitwise status, and the ratio to native RCCL. Native RCCL remains a performance +reference only, not the strict arithmetic reference. diff --git a/rl_engine/distributed/__init__.py b/rl_engine/distributed/__init__.py index 37698f1a..f83f2520 100644 --- a/rl_engine/distributed/__init__.py +++ b/rl_engine/distributed/__init__.py @@ -2,5 +2,15 @@ # Copyright (c) 2026 RL-Kernel Contributors from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.distributed.transport_collectives import ( + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, + create_deterministic_collective, +) -__all__ = ["DeterministicCollective"] +__all__ = [ + "DeterministicCollective", + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "create_deterministic_collective", +] diff --git a/rl_engine/distributed/transport_collectives.py b/rl_engine/distributed/transport_collectives.py new file mode 100644 index 00000000..fea577bc --- /dev/null +++ b/rl_engine/distributed/transport_collectives.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic reductions built on rank-ordered tensor transport. + +The transport in this module never performs a floating-point reduction. It +only gathers every rank's input, after which each rank evaluates the same +balanced reduction tree locally. In a ROCm build, PyTorch's ``nccl`` backend +is RCCL and therefore ``all_gather_into_tensor`` provides the transport. +""" + +from __future__ import annotations + +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 TorchDistributedDeterministicCollective: + """Correctness-first collectives using AllGather as transport only. + + Rank inputs are gathered without arithmetic and reduced locally as the + balanced tree ``((rank0 + rank1) + (rank2 + rank3)) + ...``. Consequently, + all ranks execute the exact same floating-point expression. TP sizes 1, + 2, 4, and 8 are nested prefixes of that expression and match the existing + CUDA IPC collective's ordering. + + The generic class also supports a CPU/Gloo process group, which is useful + as an executable reference. Production ROCm callers should use + :class:`RCCLDeterministicCollective` or + :func:`create_deterministic_collective` so backend validation fails closed. + All ranks must call methods in the same order with matching input shapes + and dtypes, and construct the instance with the same ``max_size_bytes``. + """ + + backend_id = "torch_distributed_balanced_tree" + transport_only = True + reduction_order = "balanced_rank_tree" + supports_async_overlap = False + supports_compute_communication_fusion = False + + 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 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 = int(dist.get_rank(group=self.group)) + self.world_size = int(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}" + ) + + self.device = self._normalize_device(device) + self.max_size_bytes = int(max_size_bytes) + self._backend = str(dist.get_backend(self.group)).lower() + self._lock = threading.Lock() + self._closed = False + # One dtype-agnostic byte workspace is grown on demand and reused by + # reduction collectives. AllGather writes directly into its output. + self._workspace: torch.Tensor | None = None + self._validate_matching_capacity() + + @staticmethod + def _normalize_device( + device: torch.device | str | int | None, + ) -> torch.device: + if device is None: + if torch.cuda.is_available(): + normalized = torch.device("cuda", torch.cuda.current_device()) + else: + normalized = torch.device("cpu") + elif isinstance(device, int): + normalized = torch.device("cuda", device) + else: + normalized = torch.device(device) + + if normalized.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("a CUDA/ROCm device was requested but none is available") + current_device = torch.cuda.current_device() + if normalized.index is None: + normalized = torch.device("cuda", current_device) + if normalized.index != current_device: + raise ValueError( + "the collective device must be the current CUDA/ROCm device; call " + f"torch.cuda.set_device({normalized.index}) first" + ) + return normalized + + @property + def closed(self) -> bool: + """Whether this instance rejects further collective calls.""" + + return self._closed + + @property + def workspace_size_bytes(self) -> int: + """Currently retained reduction workspace size in bytes.""" + + workspace = self._workspace + return 0 if workspace is None else int(workspace.numel()) + + def all_reduce( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the fixed balanced-tree sum on every rank.""" + + self._check_open() + self._validate_reduction_input(input) + if out is None: + out = torch.empty_like(input) + self._validate_output(out, input, tuple(input.shape)) + + with self._lock: + self._check_open() + self._validate_matching_signature("all_reduce", input) + rank_inputs = self._all_gather_transport(input) + reduced = self._balanced_tree_sum(rank_inputs) + out.copy_(reduced) + return out + + def all_gather( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Gather rank-ordered input bit patterns along dimension 0.""" + + self._check_open() + self._validate_gather_input(input) + output_shape = (input.size(0) * self.world_size, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + + with self._lock: + self._check_open() + self._validate_matching_signature("all_gather", input) + self._all_gather_transport(input, gathered_flat=out.view(-1)) + return out + + def reduce_scatter( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Fixed-tree sum followed by rank-ordered dimension-0 slicing.""" + + self._check_open() + self._validate_reduction_input(input) + if input.dim() == 0: + raise ValueError("reduce_scatter input must have at least one dimension") + if input.size(0) % self.world_size != 0: + raise ValueError( + "reduce_scatter input.size(0) must be divisible by " + f"world_size={self.world_size}; got {input.size(0)}" + ) + rows_per_rank = input.size(0) // self.world_size + output_shape = (rows_per_rank, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + + with self._lock: + self._check_open() + self._validate_matching_signature("reduce_scatter", input) + rank_inputs = self._all_gather_transport(input) + reduced = self._balanced_tree_sum(rank_inputs) + begin = self.rank * rows_per_rank + out.copy_(reduced.narrow(0, begin, rows_per_rank)) + return out + + def close(self) -> None: + """Close the instance. + + Closing releases the lazily allocated reduction workspace and marks + the lifecycle boundary. Collective calls are blocking at this API. + """ + + with self._lock: + self._workspace = None + self._closed = True + + def __enter__(self) -> TorchDistributedDeterministicCollective: + 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 getattr(self, "_closed", True): + raise RuntimeError("deterministic collective is closed") + + def _validate_tensor(self, input: torch.Tensor) -> None: + if not isinstance(input, torch.Tensor): + raise TypeError(f"input must be a torch.Tensor, got {type(input)!r}") + if 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") + 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_reduction_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dtype not in _REDUCTION_DTYPES: + raise TypeError( + "deterministic reductions support float32, float16, and bfloat16; " + f"got {input.dtype}" + ) + + def _validate_gather_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dim() == 0: + raise ValueError("all_gather input must have at least one dimension") + + @staticmethod + def _validate_output( + output: torch.Tensor, + input: torch.Tensor, + output_shape: tuple[int, ...], + ) -> None: + if not isinstance(output, torch.Tensor): + raise TypeError(f"out must be a torch.Tensor, got {type(output)!r}") + 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 tuple(output.shape) != output_shape: + raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: + if self.world_size == 1: + return + 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 _validate_matching_capacity(self) -> None: + if self.world_size == 1: + return + capacities: list[int | None] = [None] * self.world_size + dist.all_gather_object(capacities, self.max_size_bytes, group=self.group) + if any(peer_capacity != self.max_size_bytes for peer_capacity in capacities): + raise ValueError(f"all ranks must use the same max_size_bytes; got {capacities}") + + def _all_gather_transport( + self, + input: torch.Tensor, + *, + gathered_flat: torch.Tensor | None = None, + ) -> torch.Tensor: + # Flattening makes the output contract independent of whether a given + # ProcessGroup implements the concatenation or stacking form of AG. + if self.world_size == 1: + if gathered_flat is None: + gathered_flat = input.clone().view(-1) + else: + gathered_flat.copy_(input.view(-1)) + return gathered_flat.reshape((1, *input.shape)) + input_flat = input.view(-1) + required_elements = self.world_size * input_flat.numel() + if gathered_flat is None: + gathered_flat = self._workspace_for(input, required_elements) + elif ( + gathered_flat.numel() != required_elements + or gathered_flat.dtype != input.dtype + or gathered_flat.device != input.device + or not gathered_flat.is_contiguous() + ): + raise ValueError("gathered transport output has an invalid layout") + if "nccl" in self._backend: + # PyTorch exposes RCCL through the NCCL ProcessGroup API. Keep + # this as a tensor-only transport; reduction happens below. + dist.all_gather_into_tensor(gathered_flat, input_flat, group=self.group) + else: + # Some reference backends (notably Gloo versions without + # all_gather_into_tensor) only implement the list API. + gathered_chunks = list( + gathered_flat.reshape(self.world_size, input_flat.numel()).unbind(0) + ) + dist.all_gather(gathered_chunks, input_flat, group=self.group) + return gathered_flat.reshape((self.world_size, *input.shape)) + + def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.Tensor: + required_bytes = required_elements * input.element_size() + workspace = self._workspace + if workspace is None or workspace.numel() < required_bytes: + workspace = torch.empty(required_bytes, dtype=torch.uint8, device=self.device) + self._workspace = workspace + # Tensor.view(dtype) reinterprets the aligned byte allocation without + # an allocation or copy. Restrict the view to the current operation. + return workspace[:required_bytes].view(input.dtype) + + @staticmethod + def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: + level = list(rank_inputs.unbind(0)) + if len(level) not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "balanced reduction requires rank inputs for world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {len(level)}" + ) + while len(level) > 1: + level = [torch.add(level[index], level[index + 1]) for index in range(0, len(level), 2)] + return level[0] + + +class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): + """ROCm collective using RCCL AllGather strictly as tensor transport.""" + + backend_id = "rccl_all_gather_balanced_tree" + + 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 getattr(torch.version, "hip", None) is None: + raise RuntimeError("RCCL deterministic collectives require a ROCm PyTorch build") + if not torch.cuda.is_available(): + raise RuntimeError("RCCL deterministic collectives require an available ROCm device") + if device is not None and not isinstance(device, int): + requested_device = torch.device(device) + if requested_device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + super().__init__(group=group, device=device, max_size_bytes=max_size_bytes) + if self.device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + if "nccl" not in self._backend: + raise RuntimeError( + "RCCL deterministic collectives require PyTorch's NCCL process-group API" + ) + + +def create_deterministic_collective( + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, +) -> Any: + """Create the platform-appropriate deterministic collective. + + CUDA keeps using the existing IPC implementation. ROCm uses RCCL only to + gather rank inputs, followed by the shared local balanced-tree reduction. + The returned object has independent ownership; callers may cache it by + process-group/rank/device and must close an entry before replacing it. + """ + + if getattr(torch.version, "hip", None) is not None: + return RCCLDeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + # Import lazily to keep the existing CUDA implementation and its extension + # checks independent of the generic transport reference above. + from rl_engine.distributed.collectives import DeterministicCollective + + return DeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + +__all__ = [ + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "create_deterministic_collective", +] diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 91fa6f99..d5e45b96 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -27,6 +27,7 @@ CPCommunicationStatus, CUDAAGRSAttentionCPCommunication, P2PNCCLAttentionCPCommunication, + RCCLAGRSAttentionCPCommunication, sort_attention_cp_partial_states, ) except ModuleNotFoundError as exc: @@ -46,6 +47,7 @@ "CPCommunicationStatus", "CUDAAGRSAttentionCPCommunication", "P2PNCCLAttentionCPCommunication", + "RCCLAGRSAttentionCPCommunication", "sort_attention_cp_partial_states", ] diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index d1470bfd..4f945f14 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -40,7 +40,12 @@ import torch -CPCommunicationBackend = Literal["cuda_ag_rs", "p2p_nccl_reference", "local_debug"] +CPCommunicationBackend = Literal[ + "cuda_ag_rs", + "rccl_ag_rs", + "p2p_nccl_reference", + "local_debug", +] CPCommunicationStatus = Literal["interface_only", "implemented"] @@ -198,12 +203,17 @@ class AttentionCPCommunicationPlan: def validate(self) -> None: self.parallel.validate() - if self.backend not in {"cuda_ag_rs", "p2p_nccl_reference", "local_debug"}: + if self.backend not in { + "cuda_ag_rs", + "rccl_ag_rs", + "p2p_nccl_reference", + "local_debug", + }: raise ValueError(f"unsupported CP communication backend: {self.backend}") if self.status not in {"interface_only", "implemented"}: raise ValueError(f"unsupported CP communication status: {self.status}") if self.pattern != "ag_rs": - raise ValueError("PR7 CP communication must use the custom CUDA AG/RS interface") + raise ValueError("PR7 CP communication must use the self-owned AG/RS interface") if self.compute_communication != "decoupled": raise ValueError("PR7 CP communication must keep compute and communication decoupled") if self.merge_order != "global_block_index": @@ -242,6 +252,15 @@ def provenance(self) -> dict[str, object]: "cp_comm_strict_kv_communication": "all_gather", "cp_comm_strict_position_communication": "all_gather", "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_runtime": ( + "rccl" + if self.backend == "rccl_ag_rs" + else "nccl" if self.backend in {"cuda_ag_rs", "p2p_nccl_reference"} else "local" + ), + # The ROCm path intentionally transports tensors and performs the + # arithmetic in the deterministic core; only the CUDA IPC path + # owns a numeric collective reduction kernel. + "cp_comm_attention_numeric_reduction": self.backend == "cuda_ag_rs", "cp_comm_expected_kv_token_range": ( None if self.expected_kv_token_range is None else list(self.expected_kv_token_range) ), @@ -338,7 +357,10 @@ def forward( packed = full.movedim(ctx.sequence_dim, 0).contiguous() if ctx.rank != ctx.root: packed = torch.zeros_like(packed) - local = collective.reduce_scatter(packed) + # The ROCm transport adapter exposes an explicit root-owned scatter; + # CUDA's IPC collective keeps the historical reduce_scatter entrypoint. + scatter = getattr(collective, "scatter", None) + local = scatter(packed) if callable(scatter) else collective.reduce_scatter(packed) return local.movedim(0, ctx.sequence_dim).contiguous() @staticmethod @@ -557,6 +579,158 @@ def _validate_cuda_plan(self, plan: AttentionCPCommunicationPlan) -> None: raise AttentionCPCommunicationUnavailable("self-owned CUDA AG/RS requires CUDA") +class _RCCLRankOrderedTransport: + """RCCL tensor transport for the ROCm AG/RS attention contract. + + RCCL is deliberately used only for AllGather/Scatter transport. The + reduction half first gathers every source tensor and then evaluates the + same fixed balanced rank tree on every device, which avoids relying on RCCL's + implementation-defined floating-point reduction order. + """ + + def __init__(self, *, process_group: Any = None, root: int = 0) -> None: + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires initialized torch.distributed" + ) + backend = str(dist.get_backend(process_group)).lower() + if "nccl" not in backend or torch.version.hip is None: + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires the PyTorch NCCL API on ROCm" + ) + self.group = process_group + self.rank = int(dist.get_rank(process_group)) + self.world_size = int(dist.get_world_size(process_group)) + self.root = int(root) + if self.world_size not in (1, 2, 4, 8): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS supports CP world sizes 1, 2, 4, and 8" + ) + if self.root < 0 or self.root >= self.world_size: + raise AttentionCPCommunicationUnavailable("RCCL scatter root is outside the group") + + def all_gather(self, local: torch.Tensor) -> torch.Tensor: + import torch.distributed as dist + + if local.ndim == 0 or not local.is_cuda or not local.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL AllGather requires a contiguous ROCm tensor with a leading dimension" + ) + shape = (self.world_size * local.size(0), *local.shape[1:]) + gathered = torch.empty(shape, dtype=local.dtype, device=local.device) + if self.world_size == 1: + gathered.copy_(local) + else: + dist.all_gather_into_tensor(gathered, local, group=self.group) + return gathered + + def scatter(self, full: torch.Tensor) -> torch.Tensor: + import torch.distributed as dist + + if full.ndim == 0 or not full.is_cuda or not full.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL Scatter requires a contiguous ROCm tensor with a leading dimension" + ) + if full.size(0) % self.world_size: + raise AttentionCPCommunicationUnavailable( + "RCCL Scatter leading dimension must divide the CP world size" + ) + chunks = tuple(chunk.contiguous() for chunk in full.chunk(self.world_size, dim=0)) + local = torch.empty_like(chunks[self.rank]) + if self.world_size == 1: + local.copy_(chunks[0]) + return local + + # ``src`` is a global rank even when a subgroup is supplied. + global_root = self.root + if self.group is not None: + get_global_rank = getattr(dist, "get_global_rank", None) + if callable(get_global_rank): + global_root = int(get_global_rank(self.group, self.root)) + else: + get_group_ranks = getattr(dist, "get_process_group_ranks", None) + if not callable(get_group_ranks): + raise AttentionCPCommunicationUnavailable( + "PyTorch cannot map the RCCL subgroup root to a global rank" + ) + global_root = int(get_group_ranks(self.group)[self.root]) + dist.scatter( + local, + scatter_list=list(chunks) if self.rank == self.root else None, + src=global_root, + group=self.group, + ) + return local + + def reduce_scatter(self, full: torch.Tensor) -> torch.Tensor: + """Deterministic source-rank sum followed by local sequence slicing.""" + + if full.ndim == 0 or not full.is_cuda or not full.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL ReduceScatter requires a contiguous ROCm tensor with a leading dimension" + ) + if full.size(0) % self.world_size: + raise AttentionCPCommunicationUnavailable( + "RCCL ReduceScatter leading dimension must divide the CP world size" + ) + gathered = self.all_gather(full) + rows_per_rank = full.size(0) // self.world_size + begin = self.rank * rows_per_rank + # Gather layout is [source_rank, full_sequence, ...]. Evaluate the + # same balanced rank tree as the general deterministic collective. + level = [] + for source in range(self.world_size): + source_begin = source * full.size(0) + begin + level.append(gathered[source_begin : source_begin + rows_per_rank]) + while len(level) > 1: + level = [torch.add(level[index], level[index + 1]) for index in range(0, len(level), 2)] + return level[0] + + +class RCCLAGRSAttentionCPCommunication(CUDAAGRSAttentionCPCommunication): + """ROCm AG/RS adapter using RCCL only as rank-ordered tensor transport.""" + + backend_id = "rccl_ag_rs" + supports_autograd = True + transport_only = True + supports_async_overlap = False + supports_compute_communication_fusion = False + + def _get_collective(self, plan: AttentionCPCommunicationPlan): + if self._collective is None: + self._collective = _RCCLRankOrderedTransport( + process_group=self._process_group, + root=plan.merge_root_cp_rank, + ) + if self._collective.world_size != plan.parallel.cp_world_size: + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL world size does not match the CP plan" + ) + return self._collective + + def _dist(self): + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires initialized torch.distributed" + ) + return dist + + def _validate_cuda_plan(self, plan: AttentionCPCommunicationPlan) -> None: + plan.validate() + if plan.backend != "rccl_ag_rs" or plan.status != "implemented": + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires an implemented rccl_ag_rs plan" + ) + if torch.version.hip is None or not torch.cuda.is_available(): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires an available ROCm device" + ) + + class P2PNCCLAttentionCPCommunication: """Correctness-first P2P NCCL implementation of the CP protocol. @@ -1288,6 +1462,7 @@ def _rank_in_world(rank: int, world_size: int, name: str) -> None: "CPCommunicationBackend", "CPCommunicationStatus", "CUDAAGRSAttentionCPCommunication", + "RCCLAGRSAttentionCPCommunication", "P2PNCCLAttentionCPCommunication", "sort_attention_cp_partial_states", ] diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index bc0ddcef..fe17486b 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Bias-free gated FFN assembled from deterministic CUDA kernels.""" +"""Bias-free gated FFN assembled from deterministic GPU kernels.""" from __future__ import annotations @@ -35,9 +35,9 @@ def _require_ffn_kernels(*, disable_split_k: bool) -> None: if not _EXT_AVAILABLE or _C is None or missing: suffix = f" Missing symbols: {', '.join(missing)}." if missing else "" needed = ( - "compiled deterministic GEMM and SwiGLU CUDA kernels" + "compiled deterministic GEMM and SwiGLU GPU kernels" if disable_split_k - else "compiled SwiGLU CUDA kernels" + else "compiled SwiGLU GPU kernels" ) raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}") @@ -135,7 +135,8 @@ def _validate_ffn_inputs( if tensor.dtype != torch.bfloat16: raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") if not tensor.is_cuda: - raise RuntimeError(f"{name} must be on a CUDA device, got '{tensor.device}'.") + # PyTorch exposes AMD GPU tensors through the torch.cuda API too. + raise RuntimeError(f"{name} must be on a CUDA/ROCm GPU device, got '{tensor.device}'.") if tensor.device != rmsnorm_output.device: raise RuntimeError( f"all FFN inputs must be on {rmsnorm_output.device}, " @@ -143,14 +144,35 @@ def _validate_ffn_inputs( ) +def _create_collective(*, group: Any, max_size_bytes: int): + """Create the platform-specific collective lazily. + + NVIDIA keeps the existing CUDA IPC implementation. ROCm selects the + RCCL transport-only implementation, which performs the floating-point + reduction in the shared deterministic local tree. Keeping this boundary + small also makes the backend choice explicit and easy to inject in tests. + """ + + try: + from rl_engine.distributed import create_deterministic_collective + except ImportError as exc: + raise RuntimeError( + "parallel qwen3_ffn requires the platform deterministic collective " + "factory; single-device qwen3_ffn remains available" + ) from exc + + return create_deterministic_collective( + group=group, + max_size_bytes=max_size_bytes, + ) + + def _collective_for_group(group: Any, *, min_size_bytes: int): if group is None: return None import torch.distributed as dist - from rl_engine.distributed import DeterministicCollective - rank = dist.get_rank(group=group) world_size = dist.get_world_size(group=group) device_index = torch.cuda.current_device() @@ -161,7 +183,7 @@ def _collective_for_group(group: Any, *, min_size_bytes: int): if cached is not None: cached.close() - collective = DeterministicCollective( + collective = _create_collective( group=group, max_size_bytes=max(_COLLECTIVE_MIN_CAPACITY_BYTES, min_size_bytes), ) @@ -393,7 +415,8 @@ def qwen3_ffn( ``[H, I_local]``. tp_group: Optional tensor-parallel process group. Gate and Up are column-parallel; Down is row-parallel. Reductions use the - deterministic fixed-tree collectives rather than NCCL. + platform deterministic fixed-tree collectives. On ROCm, RCCL only + transports rank inputs and the reduction tree executes locally. cp_group: Optional context-parallel process group. Each rank owns different token rows and the same local weight shards. Weight gradients AllGather tokens along CP and run the full-token diff --git a/setup.py b/setup.py index 79f882d9..3f0b3d5a 100644 --- a/setup.py +++ b/setup.py @@ -139,9 +139,11 @@ def get_extensions(): "csrc/cuda/rmsnorm.cu", "csrc/cuda/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", ] if not is_rocm: + # The CUDA collective owns CUDA IPC handles and driver API calls; + # ROCm uses the Python RCCL transport implementation instead. + cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") @@ -211,9 +213,10 @@ def get_extensions(): nvcc_flags.append("-allow-unsupported-compiler") nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA" + cxx_flags = ["-O3", "-std=c++17", platform_define] extra_link_args = list(torch_rpath) - if os.name != "nt": + if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") diff --git a/tests/distributed/test_rocm_attention_transport.py b/tests/distributed/test_rocm_attention_transport.py new file mode 100644 index 00000000..327bd1f1 --- /dev/null +++ b/tests/distributed/test_rocm_attention_transport.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, + AttentionParallelSpec, + RCCLAGRSAttentionCPCommunication, +) + + +def _plan(*, backend: str = "rccl_ag_rs") -> AttentionCPCommunicationPlan: + return AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend=backend, # type: ignore[arg-type] + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + + +class _FakeRCCLTransport: + world_size = 2 + + def __init__(self) -> None: + self.gather_calls = 0 + self.scatter_calls = 0 + + def all_gather(self, tensor: torch.Tensor) -> torch.Tensor: + self.gather_calls += 1 + return torch.cat((tensor, tensor), dim=0) + + def scatter(self, tensor: torch.Tensor) -> torch.Tensor: + self.scatter_calls += 1 + return tensor.chunk(self.world_size, dim=0)[0].contiguous() + + +def test_rccl_plan_reports_transport_only_runtime() -> None: + provenance = _plan().provenance() + + assert provenance["cp_comm_runtime"] == "rccl" + assert provenance["cp_comm_attention_numeric_reduction"] is False + + +def test_rccl_adapter_uses_root_scatter_and_reports_no_fusion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + transport = _FakeRCCLTransport() + communication = RCCLAGRSAttentionCPCommunication(collective=transport) + monkeypatch.setattr(torch.version, "hip", "test", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + plan = _plan() + + local_q = torch.zeros(1, 2, 1, 4, dtype=torch.bfloat16) + assert communication.all_gather_query(local_q, plan).shape == (1, 2, 2, 4) + + full_out = torch.zeros(1, 2, 2, 4, dtype=torch.bfloat16) + full_lse = torch.zeros(1, 2, 2, dtype=torch.float32) + shard = communication.reduce_scatter_strict_result(full_out, full_lse, plan) + + assert shard.out.shape == (1, 2, 1, 4) + assert shard.lse.shape == (1, 2, 1) + assert transport.gather_calls == 1 + assert transport.scatter_calls == 2 + assert communication.transport_only is True + assert communication.supports_async_overlap is False + assert communication.supports_compute_communication_fusion is False + + +def test_rccl_adapter_fails_closed_outside_rocm(monkeypatch: pytest.MonkeyPatch) -> None: + communication = RCCLAGRSAttentionCPCommunication(collective=_FakeRCCLTransport()) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + with pytest.raises(AttentionCPCommunicationUnavailable, match="ROCm device"): + communication.all_gather_query(torch.zeros(1, 2, 1, 4), _plan()) + + +def test_rccl_adapter_rejects_cuda_plan(monkeypatch: pytest.MonkeyPatch) -> None: + communication = RCCLAGRSAttentionCPCommunication(collective=_FakeRCCLTransport()) + monkeypatch.setattr(torch.version, "hip", "test", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + with pytest.raises(AttentionCPCommunicationUnavailable, match="rccl_ag_rs plan"): + communication.all_gather_query(torch.zeros(1, 2, 1, 4), _plan(backend="cuda_ag_rs")) diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py new file mode 100644 index 00000000..00c987e1 --- /dev/null +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +import rl_engine.distributed.collectives as cuda_collectives +import rl_engine.distributed.transport_collectives as transport_collectives +from rl_engine.distributed import ( + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, +) + + +class _FakeDistributed: + """Single-process model of rank-ordered AllGather transport.""" + + def __init__( + self, + peer_inputs: list[torch.Tensor], + *, + rank: int = 0, + backend: str = "gloo", + peer_signatures: list[tuple[Any, ...]] | None = None, + peer_capacities: list[int] | None = None, + ) -> None: + self.peer_inputs = peer_inputs + self.rank = rank + self.backend = backend + self.peer_signatures = peer_signatures + self.peer_capacities = peer_capacities + self.into_tensor_calls = 0 + self.list_transport_calls = 0 + self.last_transport_output: torch.Tensor | None = None + + @property + def tensor_transport_calls(self) -> int: + return self.into_tensor_calls + self.list_transport_calls + + @staticmethod + def is_available() -> bool: + return True + + @staticmethod + def is_initialized() -> bool: + return True + + def get_rank(self, *, group: Any) -> int: + return self.rank + + def get_world_size(self, *, group: Any) -> int: + return len(self.peer_inputs) + + def get_backend(self, group: Any) -> str: + return self.backend + + def all_gather_object(self, output: list[Any], value: Any, *, group: Any) -> None: + if isinstance(value, int): + values = self.peer_capacities or [value] * len(self.peer_inputs) + else: + values = self.peer_signatures or [value] * len(self.peer_inputs) + output[:] = values + + def all_gather_into_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + *, + group: Any, + ) -> None: + self.into_tensor_calls += 1 + self.last_transport_output = output + gathered = torch.cat([peer.reshape(-1) for peer in self.peer_inputs]) + output.copy_(gathered) + + def all_gather( + self, + output: list[torch.Tensor], + input: torch.Tensor, + *, + group: Any, + ) -> None: + self.list_transport_calls += 1 + self.last_transport_output = output[0]._base + for destination, peer in zip(output, self.peer_inputs, strict=True): + destination.copy_(peer.reshape(-1)) + + +def _make_collective( + monkeypatch: pytest.MonkeyPatch, + peer_inputs: list[torch.Tensor], + *, + rank: int = 0, + backend: str = "gloo", + peer_signatures: list[tuple[Any, ...]] | None = None, + peer_capacities: list[int] | None = None, + max_size_bytes: int = 1024, +) -> tuple[TorchDistributedDeterministicCollective, _FakeDistributed]: + fake_dist = _FakeDistributed( + peer_inputs, + rank=rank, + backend=backend, + peer_signatures=peer_signatures, + peer_capacities=peer_capacities, + ) + monkeypatch.setattr(transport_collectives, "dist", fake_dist) + collective = TorchDistributedDeterministicCollective( + group=object(), + device="cpu", + max_size_bytes=max_size_bytes, + ) + return collective, fake_dist + + +@pytest.mark.parametrize("world_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_all_reduce_supports_fixed_world_sizes_and_dtypes( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + dtype: torch.dtype, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=dtype) for rank in range(world_size)] + collective, fake_dist = _make_collective(monkeypatch, peers) + + provided = torch.empty_like(peers[0]) + returned = collective.all_reduce(peers[0], out=provided) + + assert returned is provided + assert torch.equal(provided, torch.full_like(provided, world_size * (world_size + 1) // 2)) + assert fake_dist.tensor_transport_calls == (0 if world_size == 1 else 1) + + +def test_all_reduce_uses_balanced_not_rank_ordered_left_fold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Balanced: (1e20 + 1) + (-1e20 + 1) == 0 in FP32. + # Left fold: ((1e20 + 1) + -1e20) + 1 == 1 in FP32. + peers = [torch.tensor([value], dtype=torch.float32) for value in (1.0e20, 1.0, -1.0e20, 1.0)] + collective, _ = _make_collective(monkeypatch, peers) + + output = collective.all_reduce(peers[0]) + + assert torch.equal(output, torch.zeros_like(output)) + + +def test_all_reduce_stages_before_writing_in_place(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(4)] + collective, _ = _make_collective(monkeypatch, peers) + local = peers[0].clone() + + returned = collective.all_reduce(local, out=local) + + assert returned is local + assert torch.equal(local, torch.full_like(local, 10)) + + +def test_all_gather_is_rank_ordered_and_transport_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.arange(rank * 6, (rank + 1) * 6).reshape(2, 3) for rank in range(4)] + collective, fake_dist = _make_collective(monkeypatch, peers, rank=2) + + output = collective.all_gather(peers[2]) + + assert torch.equal(output, torch.cat(peers, dim=0)) + assert fake_dist.tensor_transport_calls == 1 + + +def test_nccl_backend_uses_all_gather_into_tensor_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + + output = collective.all_gather(peers[0]) + + assert torch.equal(output, torch.cat(peers, dim=0)) + assert fake_dist.into_tensor_calls == 1 + assert fake_dist.list_transport_calls == 0 + + +def test_all_gather_writes_directly_to_provided_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + provided = torch.empty(4, 3, dtype=peers[0].dtype) + + returned = collective.all_gather(peers[0], out=provided) + + assert returned is provided + assert fake_dist.last_transport_output is not None + assert fake_dist.last_transport_output.data_ptr() == provided.data_ptr() + assert collective._workspace is None + + +def test_reduction_workspace_grows_once_and_is_reused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(2)] + collective, _ = _make_collective(monkeypatch, peers, backend="nccl") + + collective.all_reduce(peers[0]) + assert collective._workspace is not None + assert collective.workspace_size_bytes == sum(peer.numel() for peer in peers) * 4 + first_pointer = collective._workspace.data_ptr() + collective.all_reduce(peers[0]) + + assert collective._workspace.data_ptr() == first_pointer + collective.close() + assert collective._workspace is None + assert collective.workspace_size_bytes == 0 + + +def test_reduce_scatter_reduces_then_selects_local_leading_shard( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((8, 3), rank + 1, dtype=torch.bfloat16) for rank in range(4)] + collective, _ = _make_collective(monkeypatch, peers, rank=2) + + output = collective.reduce_scatter(peers[2]) + + expected_full = torch.full_like(peers[0], 10) + assert torch.equal(output, expected_full.chunk(4, dim=0)[2]) + + +def test_matching_signature_is_checked_before_tensor_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(2, 3), torch.ones(2, 3)] + signatures = [ + ("all_reduce", (2, 3), "torch.float32", 6), + ("reduce_scatter", (2, 3), "torch.float32", 6), + ] + collective, fake_dist = _make_collective( + monkeypatch, + peers, + peer_signatures=signatures, + ) + + with pytest.raises(ValueError, match="matching shapes and dtypes"): + collective.all_reduce(peers[0]) + assert fake_dist.tensor_transport_calls == 0 + + +def test_capacity_must_match_on_every_rank(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.ones(2, 3), torch.ones(2, 3)] + + with pytest.raises(ValueError, match="same max_size_bytes"): + _make_collective( + monkeypatch, + peers, + max_size_bytes=1024, + peer_capacities=[1024, 2048], + ) + + +def test_validation_fails_closed_before_transport(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.ones(4, 2), torch.ones(4, 2)] + collective, fake_dist = _make_collective(monkeypatch, peers, max_size_bytes=31) + + with pytest.raises(TypeError, match="float32, float16, and bfloat16"): + collective.all_reduce(torch.ones(1, dtype=torch.int32)) + with pytest.raises(ValueError, match="max_size_bytes"): + collective.all_reduce(peers[0]) + with pytest.raises(ValueError, match=r"input.size\(0\).+divisible"): + collective.reduce_scatter(torch.ones(3, 2)) + with pytest.raises(ValueError, match="at least one dimension"): + collective.all_gather(torch.tensor(1.0)) + assert fake_dist.tensor_transport_calls == 0 + + +def test_lifecycle_is_idempotent_and_context_manager_closes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(2, 3)] + collective, _ = _make_collective(monkeypatch, peers) + + with collective as entered: + assert entered is collective + assert not collective.closed + assert torch.equal(collective.all_reduce(peers[0]), peers[0]) + + assert collective.closed + collective.close() + with pytest.raises(RuntimeError, match="closed"): + collective.all_reduce(peers[0]) + + +@pytest.mark.parametrize("world_size", [3, 16]) +def test_unsupported_world_size_is_rejected( + monkeypatch: pytest.MonkeyPatch, + world_size: int, +) -> None: + fake_dist = _FakeDistributed([torch.ones(1)] * world_size) + monkeypatch.setattr(transport_collectives, "dist", fake_dist) + + with pytest.raises(ValueError, match="world_size in"): + TorchDistributedDeterministicCollective(group=object(), device="cpu") + + +def test_rccl_class_requires_rocm_build(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(transport_collectives.torch.version, "hip", None, raising=False) + + with pytest.raises(RuntimeError, match="ROCm PyTorch build"): + RCCLDeterministicCollective(group=object(), device="cuda:0") + + +def test_rccl_class_requires_nccl_process_group(monkeypatch: pytest.MonkeyPatch) -> None: + fake_dist = _FakeDistributed([torch.ones(1)], backend="gloo") + monkeypatch.setattr(transport_collectives, "dist", fake_dist) + monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(transport_collectives.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(transport_collectives.torch.cuda, "current_device", lambda: 0) + + with pytest.raises(RuntimeError, match="NCCL process-group API"): + RCCLDeterministicCollective(group=object(), device="cuda:0") + + +def test_rccl_class_rejects_cpu_before_process_group_exchange( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(transport_collectives.torch.cuda, "is_available", lambda: True) + + with pytest.raises(ValueError, match="ROCm device"): + RCCLDeterministicCollective(group=object(), device="cpu") + + +def test_factory_dispatches_rocm_to_rccl(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + calls: list[dict[str, Any]] = [] + + def fake_rccl(**kwargs: Any) -> object: + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(transport_collectives, "RCCLDeterministicCollective", fake_rccl) + + group = object() + result = transport_collectives.create_deterministic_collective( + group=group, + device="cuda:3", + max_size_bytes=1234, + ) + + assert result is sentinel + assert calls == [{"group": group, "device": "cuda:3", "max_size_bytes": 1234}] + + +def test_factory_preserves_existing_cuda_collective(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + calls: list[dict[str, Any]] = [] + + def fake_cuda(**kwargs: Any) -> object: + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(transport_collectives.torch.version, "hip", None, raising=False) + monkeypatch.setattr(cuda_collectives, "DeterministicCollective", fake_cuda) + + group = object() + result = transport_collectives.create_deterministic_collective( + group=group, + device="cuda:1", + max_size_bytes=4321, + ) + + assert result is sentinel + assert calls == [{"group": group, "device": "cuda:1", "max_size_bytes": 4321}] diff --git a/tests/test_build_platform_collectives.py b/tests/test_build_platform_collectives.py new file mode 100644 index 00000000..f094a5af --- /dev/null +++ b/tests/test_build_platform_collectives.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import runpy +from typing import Any + +import setuptools +import torch +from torch.utils import cpp_extension + + +def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def fake_setup(**kwargs: Any) -> None: + captured.update(kwargs) + + def fake_extension(**kwargs: Any) -> dict[str, Any]: + return kwargs + + monkeypatch.setattr(setuptools, "setup", fake_setup) + monkeypatch.setattr(cpp_extension, "CUDAExtension", fake_extension) + monkeypatch.setattr(torch.version, "hip", hip, raising=False) + monkeypatch.delenv("KERNEL_ALIGN_FORCE_SM90", raising=False) + monkeypatch.delenv("KERNEL_ALIGN_DET_GEMM_SM90", raising=False) + if hip is None: + monkeypatch.delenv("PYTORCH_ROCM_ARCH", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0)) + else: + monkeypatch.setenv("PYTORCH_ROCM_ARCH", "gfx942") + + runpy.run_path("setup.py", run_name=f"rl_kernel_setup_probe_{hip or 'cuda'}") + return captured["ext_modules"][0] + + +def test_rocm_build_excludes_cuda_ipc_collective_and_driver(monkeypatch) -> None: + extension = _load_extension_config(monkeypatch, hip="test") + + assert "csrc/cuda/distributed/deterministic_collective.cu" not in extension["sources"] + assert "-DKERNEL_ALIGN_WITH_ROCM" in extension["extra_compile_args"]["cxx"] + assert "-DKERNEL_ALIGN_WITH_CUDA" not in extension["extra_compile_args"]["cxx"] + assert "-lcuda" not in extension["extra_link_args"] + + +def test_cuda_build_keeps_existing_ipc_collective(monkeypatch) -> None: + extension = _load_extension_config(monkeypatch, hip=None) + + assert "csrc/cuda/distributed/deterministic_collective.cu" in extension["sources"] + assert "-DKERNEL_ALIGN_WITH_CUDA" in extension["extra_compile_args"]["cxx"] + assert "-DKERNEL_ALIGN_WITH_ROCM" not in extension["extra_compile_args"]["cxx"] + assert "-lcuda" in extension["extra_link_args"] diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index 515abc7a..fddfa9f4 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -137,6 +137,25 @@ def _close_ffn_collectives() -> None: ffn_module._COLLECTIVES.clear() +def test_ffn_collective_creation_uses_platform_factory(monkeypatch): + import rl_engine.distributed as distributed + + sentinel = object() + calls = [] + + def fake_factory(**kwargs): + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(distributed, "create_deterministic_collective", fake_factory) + group = object() + + result = ffn_module._create_collective(group=group, max_size_bytes=1234) + + assert result is sentinel + assert calls == [{"group": group, "max_size_bytes": 1234}] + + def _shard_ranges( rank: int, *, diff --git a/tests/test_rocm_collective_benchmark.py b/tests/test_rocm_collective_benchmark.py new file mode 100644 index 00000000..e4eb8643 --- /dev/null +++ b/tests/test_rocm_collective_benchmark.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + +_BENCHMARK_PATH = Path(__file__).parents[1] / "benchmarks" / "benchmark_rocm_collectives.py" +_SPEC = importlib.util.spec_from_file_location( + "rlkernel_rocm_collective_benchmark", _BENCHMARK_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +benchmark = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(benchmark) + + +def test_rocm_collective_benchmark_parser() -> None: + args = benchmark.parse_args( + [ + "--size-bytes", + "1024", + "4096", + "--dtype", + "fp32", + "--operations", + "all_reduce", + "reduce_scatter", + "--warmup", + "2", + "--iterations", + "3", + "--samples", + "4", + ] + ) + + benchmark._validate_args(args) + assert args.size_bytes == [1024, 4096] + assert args.dtype == "fp32" + assert args.operations == ["all_reduce", "reduce_scatter"] + assert (args.warmup, args.iterations, args.samples) == (2, 3, 4) + + +@pytest.mark.parametrize( + "argv", + ( + ["--size-bytes", "0"], + ["--warmup", "-1"], + ["--iterations", "0"], + ["--samples", "0"], + ), +) +def test_rocm_collective_benchmark_rejects_invalid_counts(argv: list[str]) -> None: + with pytest.raises(ValueError): + benchmark._validate_args(benchmark.parse_args(argv)) + + +def test_rocm_collective_benchmark_aligns_reduce_scatter_input() -> None: + tensor, actual_bytes = benchmark._make_inputs( + size_bytes=35, + dtype=torch.float32, + world_size=4, + rank=2, + device=torch.device("cpu"), + ) + + assert tensor.is_contiguous() + assert tensor.numel() % 4 == 0 + assert actual_bytes == tensor.numel() * tensor.element_size() From eb16e79bbfecdb975786adf95c1015f5e60f4ff7 Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 16:48:51 +0000 Subject: [PATCH 2/5] perf(distributed): remove ROCm collective hot-path allocations --- .../distributed/transport_collectives.py | 28 +++++++++++++++---- .../kernels/ops/cuda/attention/cp_comm.py | 7 +++-- ...test_transport_deterministic_collective.py | 15 ++++++++++ 3 files changed, 42 insertions(+), 8 deletions(-) diff --git a/rl_engine/distributed/transport_collectives.py b/rl_engine/distributed/transport_collectives.py index fea577bc..ff9a4dcf 100644 --- a/rl_engine/distributed/transport_collectives.py +++ b/rl_engine/distributed/transport_collectives.py @@ -74,6 +74,12 @@ def __init__( # One dtype-agnostic byte workspace is grown on demand and reused by # reduction collectives. AllGather writes directly into its output. self._workspace: torch.Tensor | None = None + # A Python-object collective is useful for catching a mismatched new + # signature, but running one on every hot-path call dominates small + # message latency. Validate each local signature once and then rely on + # the standard collective contract that ranks call operations in the + # same order. + self._validated_signatures: set[tuple[Any, ...]] = set() self._validate_matching_capacity() @staticmethod @@ -200,6 +206,7 @@ def close(self) -> None: with self._lock: self._workspace = None + self._validated_signatures.clear() self._closed = True def __enter__(self) -> TorchDistributedDeterministicCollective: @@ -271,12 +278,15 @@ def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> Non if self.world_size == 1: return signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + if signature in self._validated_signatures: + return 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}" ) + self._validated_signatures.add(signature) def _validate_matching_capacity(self) -> None: if self.world_size == 1: @@ -336,15 +346,21 @@ def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.T @staticmethod def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: - level = list(rank_inputs.unbind(0)) - if len(level) not in _SUPPORTED_WORLD_SIZES: + world_size = rank_inputs.size(0) + if world_size not in _SUPPORTED_WORLD_SIZES: raise ValueError( "balanced reduction requires rank inputs for world_size in " - f"{_SUPPORTED_WORLD_SIZES}, got {len(level)}" + f"{_SUPPORTED_WORLD_SIZES}, got {world_size}" ) - while len(level) > 1: - level = [torch.add(level[index], level[index + 1]) for index in range(0, len(level), 2)] - return level[0] + # ``rank_inputs`` is the private transport workspace for reductions, + # so fixed-tree nodes can be accumulated in place. This preserves the + # exact pairings while avoiding one temporary allocation per tree node. + stride = 1 + while stride < world_size: + for index in range(0, world_size, 2 * stride): + rank_inputs[index].add_(rank_inputs[index + stride]) + stride *= 2 + return rank_inputs[0] class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index 4f945f14..9f805fe5 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -684,8 +684,11 @@ def reduce_scatter(self, full: torch.Tensor) -> torch.Tensor: for source in range(self.world_size): source_begin = source * full.size(0) + begin level.append(gathered[source_begin : source_begin + rows_per_rank]) - while len(level) > 1: - level = [torch.add(level[index], level[index + 1]) for index in range(0, len(level), 2)] + stride = 1 + while stride < self.world_size: + for index in range(0, self.world_size, 2 * stride): + level[index].add_(level[index + stride]) + stride *= 2 return level[0] diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py index 00c987e1..ed1da194 100644 --- a/tests/distributed/test_transport_deterministic_collective.py +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -35,6 +35,7 @@ def __init__( self.peer_capacities = peer_capacities self.into_tensor_calls = 0 self.list_transport_calls = 0 + self.object_gather_calls = 0 self.last_transport_output: torch.Tensor | None = None @property @@ -59,6 +60,7 @@ def get_backend(self, group: Any) -> str: return self.backend def all_gather_object(self, output: list[Any], value: Any, *, group: Any) -> None: + self.object_gather_calls += 1 if isinstance(value, int): values = self.peer_capacities or [value] * len(self.peer_inputs) else: @@ -216,6 +218,19 @@ def test_reduction_workspace_grows_once_and_is_reused( assert collective.workspace_size_bytes == 0 +def test_matching_signature_is_validated_once_per_hot_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + constructor_object_gathers = fake_dist.object_gather_calls + + collective.all_reduce(peers[0]) + collective.all_reduce(peers[0]) + + assert fake_dist.object_gather_calls == constructor_object_gathers + 1 + + def test_reduce_scatter_reduces_then_selects_local_leading_shard( monkeypatch: pytest.MonkeyPatch, ) -> None: From d36d70bb0c5ed1ec72ad563aff5ef834af85a38f Mon Sep 17 00:00:00 2001 From: vensen Date: Fri, 28 Aug 2026 17:37:49 +0000 Subject: [PATCH 3/5] refactor(distributed): unify platform collectives module --- csrc/ops.cpp | 2 +- rl_engine/distributed/__init__.py | 4 +- rl_engine/distributed/collectives.py | 416 +++++++++++++++++ .../distributed/transport_collectives.py | 435 ------------------ ...test_transport_deterministic_collective.py | 46 +- 5 files changed, 448 insertions(+), 455 deletions(-) delete mode 100644 rl_engine/distributed/transport_collectives.py diff --git a/csrc/ops.cpp b/csrc/ops.cpp index fb84dc64..8936c9fb 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -95,7 +95,7 @@ torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torc #if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) // Single-node TP=8 deterministic CUDA IPC collectives. ROCm uses the -// rank-ordered RCCL transport in rl_engine.distributed.transport_collectives. +// rank-ordered RCCL transport in rl_engine.distributed.collectives. std::tuple, int64_t> deterministic_collective_ipc_meta( torch::Tensor& tensor); int64_t deterministic_collective_create( diff --git a/rl_engine/distributed/__init__.py b/rl_engine/distributed/__init__.py index f83f2520..9010a534 100644 --- a/rl_engine/distributed/__init__.py +++ b/rl_engine/distributed/__init__.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from rl_engine.distributed.collectives import DeterministicCollective -from rl_engine.distributed.transport_collectives import ( +from rl_engine.distributed.collectives import ( + DeterministicCollective, RCCLDeterministicCollective, TorchDistributedDeterministicCollective, create_deterministic_collective, diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 54b00fa5..db3ac70f 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -1,5 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic collectives for native CUDA IPC and rank-ordered transport. + +The ROCm transport never performs a floating-point reduction. It gathers every +rank's input through RCCL, after which each rank evaluates the same balanced +reduction tree locally. +""" from __future__ import annotations @@ -323,3 +329,413 @@ def _synchronize_ranks(self) -> None: dist.barrier(group=self.group, device_ids=[self.device.index]) else: dist.barrier(group=self.group) + + +class TorchDistributedDeterministicCollective: + """Correctness-first collectives using AllGather as transport only. + + Rank inputs are gathered without arithmetic and reduced locally as the + balanced tree ``((rank0 + rank1) + (rank2 + rank3)) + ...``. Consequently, + all ranks execute the exact same floating-point expression. TP sizes 1, + 2, 4, and 8 are nested prefixes of that expression and match the existing + CUDA IPC collective's ordering. + + The generic class also supports a CPU/Gloo process group, which is useful + as an executable reference. Production ROCm callers should use + :class:`RCCLDeterministicCollective` or + :func:`create_deterministic_collective` so backend validation fails closed. + All ranks must call methods in the same order with matching input shapes + and dtypes, and construct the instance with the same ``max_size_bytes``. + """ + + backend_id = "torch_distributed_balanced_tree" + transport_only = True + reduction_order = "balanced_rank_tree" + supports_async_overlap = False + supports_compute_communication_fusion = False + + 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 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 = int(dist.get_rank(group=self.group)) + self.world_size = int(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}" + ) + + self.device = self._normalize_device(device) + self.max_size_bytes = int(max_size_bytes) + self._backend = str(dist.get_backend(self.group)).lower() + self._lock = threading.Lock() + self._closed = False + # One dtype-agnostic byte workspace is grown on demand and reused by + # reduction collectives. AllGather writes directly into its output. + self._workspace: torch.Tensor | None = None + # A Python-object collective is useful for catching a mismatched new + # signature, but running one on every hot-path call dominates small + # message latency. Validate each local signature once and then rely on + # the standard collective contract that ranks call operations in the + # same order. + self._validated_signatures: set[tuple[Any, ...]] = set() + self._validate_matching_capacity() + + @staticmethod + def _normalize_device( + device: torch.device | str | int | None, + ) -> torch.device: + if device is None: + if torch.cuda.is_available(): + normalized = torch.device("cuda", torch.cuda.current_device()) + else: + normalized = torch.device("cpu") + elif isinstance(device, int): + normalized = torch.device("cuda", device) + else: + normalized = torch.device(device) + + if normalized.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("a CUDA/ROCm device was requested but none is available") + current_device = torch.cuda.current_device() + if normalized.index is None: + normalized = torch.device("cuda", current_device) + if normalized.index != current_device: + raise ValueError( + "the collective device must be the current CUDA/ROCm device; call " + f"torch.cuda.set_device({normalized.index}) first" + ) + return normalized + + @property + def closed(self) -> bool: + """Whether this instance rejects further collective calls.""" + + return self._closed + + @property + def workspace_size_bytes(self) -> int: + """Currently retained reduction workspace size in bytes.""" + + workspace = self._workspace + return 0 if workspace is None else int(workspace.numel()) + + def all_reduce( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the fixed balanced-tree sum on every rank.""" + + self._check_open() + self._validate_reduction_input(input) + if out is None: + out = torch.empty_like(input) + self._validate_output(out, input, tuple(input.shape)) + + with self._lock: + self._check_open() + self._validate_matching_signature("all_reduce", input) + rank_inputs = self._all_gather_transport(input) + reduced = self._balanced_tree_sum(rank_inputs) + out.copy_(reduced) + return out + + def all_gather( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Gather rank-ordered input bit patterns along dimension 0.""" + + self._check_open() + self._validate_gather_input(input) + output_shape = (input.size(0) * self.world_size, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + + with self._lock: + self._check_open() + self._validate_matching_signature("all_gather", input) + self._all_gather_transport(input, gathered_flat=out.view(-1)) + return out + + def reduce_scatter( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Fixed-tree sum followed by rank-ordered dimension-0 slicing.""" + + self._check_open() + self._validate_reduction_input(input) + if input.dim() == 0: + raise ValueError("reduce_scatter input must have at least one dimension") + if input.size(0) % self.world_size != 0: + raise ValueError( + "reduce_scatter input.size(0) must be divisible by " + f"world_size={self.world_size}; got {input.size(0)}" + ) + rows_per_rank = input.size(0) // self.world_size + output_shape = (rows_per_rank, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + + with self._lock: + self._check_open() + self._validate_matching_signature("reduce_scatter", input) + rank_inputs = self._all_gather_transport(input) + reduced = self._balanced_tree_sum(rank_inputs) + begin = self.rank * rows_per_rank + out.copy_(reduced.narrow(0, begin, rows_per_rank)) + return out + + def close(self) -> None: + """Close the instance. + + Closing releases the lazily allocated reduction workspace and marks + the lifecycle boundary. Collective calls are blocking at this API. + """ + + with self._lock: + self._workspace = None + self._validated_signatures.clear() + self._closed = True + + def __enter__(self) -> TorchDistributedDeterministicCollective: + 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 getattr(self, "_closed", True): + raise RuntimeError("deterministic collective is closed") + + def _validate_tensor(self, input: torch.Tensor) -> None: + if not isinstance(input, torch.Tensor): + raise TypeError(f"input must be a torch.Tensor, got {type(input)!r}") + if 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") + 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_reduction_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dtype not in _REDUCTION_DTYPES: + raise TypeError( + "deterministic reductions support float32, float16, and bfloat16; " + f"got {input.dtype}" + ) + + def _validate_gather_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dim() == 0: + raise ValueError("all_gather input must have at least one dimension") + + @staticmethod + def _validate_output( + output: torch.Tensor, + input: torch.Tensor, + output_shape: tuple[int, ...], + ) -> None: + if not isinstance(output, torch.Tensor): + raise TypeError(f"out must be a torch.Tensor, got {type(output)!r}") + 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 tuple(output.shape) != output_shape: + raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: + if self.world_size == 1: + return + signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + if signature in self._validated_signatures: + return + 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}" + ) + self._validated_signatures.add(signature) + + def _validate_matching_capacity(self) -> None: + if self.world_size == 1: + return + capacities: list[int | None] = [None] * self.world_size + dist.all_gather_object(capacities, self.max_size_bytes, group=self.group) + if any(peer_capacity != self.max_size_bytes for peer_capacity in capacities): + raise ValueError(f"all ranks must use the same max_size_bytes; got {capacities}") + + def _all_gather_transport( + self, + input: torch.Tensor, + *, + gathered_flat: torch.Tensor | None = None, + ) -> torch.Tensor: + # Flattening makes the output contract independent of whether a given + # ProcessGroup implements the concatenation or stacking form of AG. + if self.world_size == 1: + if gathered_flat is None: + gathered_flat = input.clone().view(-1) + else: + gathered_flat.copy_(input.view(-1)) + return gathered_flat.reshape((1, *input.shape)) + input_flat = input.view(-1) + required_elements = self.world_size * input_flat.numel() + if gathered_flat is None: + gathered_flat = self._workspace_for(input, required_elements) + elif ( + gathered_flat.numel() != required_elements + or gathered_flat.dtype != input.dtype + or gathered_flat.device != input.device + or not gathered_flat.is_contiguous() + ): + raise ValueError("gathered transport output has an invalid layout") + if "nccl" in self._backend: + # PyTorch exposes RCCL through the NCCL ProcessGroup API. Keep this + # as a tensor-only transport; reduction happens below. + dist.all_gather_into_tensor(gathered_flat, input_flat, group=self.group) + else: + # Some reference backends (notably Gloo versions without + # all_gather_into_tensor) only implement the list API. + gathered_chunks = list( + gathered_flat.reshape(self.world_size, input_flat.numel()).unbind(0) + ) + dist.all_gather(gathered_chunks, input_flat, group=self.group) + return gathered_flat.reshape((self.world_size, *input.shape)) + + def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.Tensor: + required_bytes = required_elements * input.element_size() + workspace = self._workspace + if workspace is None or workspace.numel() < required_bytes: + workspace = torch.empty(required_bytes, dtype=torch.uint8, device=self.device) + self._workspace = workspace + # Tensor.view(dtype) reinterprets the aligned byte allocation without + # an allocation or copy. Restrict the view to the current operation. + return workspace[:required_bytes].view(input.dtype) + + @staticmethod + def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: + world_size = rank_inputs.size(0) + if world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "balanced reduction requires rank inputs for world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {world_size}" + ) + # ``rank_inputs`` is the private transport workspace for reductions, + # so fixed-tree nodes can be accumulated in place. This preserves the + # exact pairings while avoiding one temporary allocation per tree node. + stride = 1 + while stride < world_size: + for index in range(0, world_size, 2 * stride): + rank_inputs[index].add_(rank_inputs[index + stride]) + stride *= 2 + return rank_inputs[0] + + +class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): + """ROCm collective using RCCL AllGather strictly as tensor transport.""" + + backend_id = "rccl_all_gather_balanced_tree" + + 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 getattr(torch.version, "hip", None) is None: + raise RuntimeError("RCCL deterministic collectives require a ROCm PyTorch build") + if not torch.cuda.is_available(): + raise RuntimeError("RCCL deterministic collectives require an available ROCm device") + if device is not None and not isinstance(device, int): + requested_device = torch.device(device) + if requested_device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + super().__init__(group=group, device=device, max_size_bytes=max_size_bytes) + if self.device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + if "nccl" not in self._backend: + raise RuntimeError( + "RCCL deterministic collectives require PyTorch's NCCL process-group API" + ) + + +def create_deterministic_collective( + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, +) -> Any: + """Create the platform-appropriate deterministic collective. + + CUDA uses the native ``DeterministicCollective`` implementation. ROCm uses + RCCL only to gather rank inputs, followed by the shared local balanced-tree + reduction. The returned object has independent ownership; callers may cache + it and must close an entry before replacing it. + """ + + if getattr(torch.version, "hip", None) is not None: + return RCCLDeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + return DeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + +__all__ = [ + "DeterministicCollective", + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "create_deterministic_collective", +] diff --git a/rl_engine/distributed/transport_collectives.py b/rl_engine/distributed/transport_collectives.py deleted file mode 100644 index ff9a4dcf..00000000 --- a/rl_engine/distributed/transport_collectives.py +++ /dev/null @@ -1,435 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic reductions built on rank-ordered tensor transport. - -The transport in this module never performs a floating-point reduction. It -only gathers every rank's input, after which each rank evaluates the same -balanced reduction tree locally. In a ROCm build, PyTorch's ``nccl`` backend -is RCCL and therefore ``all_gather_into_tensor`` provides the transport. -""" - -from __future__ import annotations - -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 TorchDistributedDeterministicCollective: - """Correctness-first collectives using AllGather as transport only. - - Rank inputs are gathered without arithmetic and reduced locally as the - balanced tree ``((rank0 + rank1) + (rank2 + rank3)) + ...``. Consequently, - all ranks execute the exact same floating-point expression. TP sizes 1, - 2, 4, and 8 are nested prefixes of that expression and match the existing - CUDA IPC collective's ordering. - - The generic class also supports a CPU/Gloo process group, which is useful - as an executable reference. Production ROCm callers should use - :class:`RCCLDeterministicCollective` or - :func:`create_deterministic_collective` so backend validation fails closed. - All ranks must call methods in the same order with matching input shapes - and dtypes, and construct the instance with the same ``max_size_bytes``. - """ - - backend_id = "torch_distributed_balanced_tree" - transport_only = True - reduction_order = "balanced_rank_tree" - supports_async_overlap = False - supports_compute_communication_fusion = False - - 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 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 = int(dist.get_rank(group=self.group)) - self.world_size = int(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}" - ) - - self.device = self._normalize_device(device) - self.max_size_bytes = int(max_size_bytes) - self._backend = str(dist.get_backend(self.group)).lower() - self._lock = threading.Lock() - self._closed = False - # One dtype-agnostic byte workspace is grown on demand and reused by - # reduction collectives. AllGather writes directly into its output. - self._workspace: torch.Tensor | None = None - # A Python-object collective is useful for catching a mismatched new - # signature, but running one on every hot-path call dominates small - # message latency. Validate each local signature once and then rely on - # the standard collective contract that ranks call operations in the - # same order. - self._validated_signatures: set[tuple[Any, ...]] = set() - self._validate_matching_capacity() - - @staticmethod - def _normalize_device( - device: torch.device | str | int | None, - ) -> torch.device: - if device is None: - if torch.cuda.is_available(): - normalized = torch.device("cuda", torch.cuda.current_device()) - else: - normalized = torch.device("cpu") - elif isinstance(device, int): - normalized = torch.device("cuda", device) - else: - normalized = torch.device(device) - - if normalized.type == "cuda": - if not torch.cuda.is_available(): - raise RuntimeError("a CUDA/ROCm device was requested but none is available") - current_device = torch.cuda.current_device() - if normalized.index is None: - normalized = torch.device("cuda", current_device) - if normalized.index != current_device: - raise ValueError( - "the collective device must be the current CUDA/ROCm device; call " - f"torch.cuda.set_device({normalized.index}) first" - ) - return normalized - - @property - def closed(self) -> bool: - """Whether this instance rejects further collective calls.""" - - return self._closed - - @property - def workspace_size_bytes(self) -> int: - """Currently retained reduction workspace size in bytes.""" - - workspace = self._workspace - return 0 if workspace is None else int(workspace.numel()) - - def all_reduce( - self, - input: torch.Tensor, - *, - out: torch.Tensor | None = None, - ) -> torch.Tensor: - """Return the fixed balanced-tree sum on every rank.""" - - self._check_open() - self._validate_reduction_input(input) - if out is None: - out = torch.empty_like(input) - self._validate_output(out, input, tuple(input.shape)) - - with self._lock: - self._check_open() - self._validate_matching_signature("all_reduce", input) - rank_inputs = self._all_gather_transport(input) - reduced = self._balanced_tree_sum(rank_inputs) - out.copy_(reduced) - return out - - def all_gather( - self, - input: torch.Tensor, - *, - out: torch.Tensor | None = None, - ) -> torch.Tensor: - """Gather rank-ordered input bit patterns along dimension 0.""" - - self._check_open() - self._validate_gather_input(input) - output_shape = (input.size(0) * self.world_size, *input.shape[1:]) - if out is None: - out = torch.empty(output_shape, dtype=input.dtype, device=input.device) - self._validate_output(out, input, output_shape) - - with self._lock: - self._check_open() - self._validate_matching_signature("all_gather", input) - self._all_gather_transport(input, gathered_flat=out.view(-1)) - return out - - def reduce_scatter( - self, - input: torch.Tensor, - *, - out: torch.Tensor | None = None, - ) -> torch.Tensor: - """Fixed-tree sum followed by rank-ordered dimension-0 slicing.""" - - self._check_open() - self._validate_reduction_input(input) - if input.dim() == 0: - raise ValueError("reduce_scatter input must have at least one dimension") - if input.size(0) % self.world_size != 0: - raise ValueError( - "reduce_scatter input.size(0) must be divisible by " - f"world_size={self.world_size}; got {input.size(0)}" - ) - rows_per_rank = input.size(0) // self.world_size - output_shape = (rows_per_rank, *input.shape[1:]) - if out is None: - out = torch.empty(output_shape, dtype=input.dtype, device=input.device) - self._validate_output(out, input, output_shape) - - with self._lock: - self._check_open() - self._validate_matching_signature("reduce_scatter", input) - rank_inputs = self._all_gather_transport(input) - reduced = self._balanced_tree_sum(rank_inputs) - begin = self.rank * rows_per_rank - out.copy_(reduced.narrow(0, begin, rows_per_rank)) - return out - - def close(self) -> None: - """Close the instance. - - Closing releases the lazily allocated reduction workspace and marks - the lifecycle boundary. Collective calls are blocking at this API. - """ - - with self._lock: - self._workspace = None - self._validated_signatures.clear() - self._closed = True - - def __enter__(self) -> TorchDistributedDeterministicCollective: - 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 getattr(self, "_closed", True): - raise RuntimeError("deterministic collective is closed") - - def _validate_tensor(self, input: torch.Tensor) -> None: - if not isinstance(input, torch.Tensor): - raise TypeError(f"input must be a torch.Tensor, got {type(input)!r}") - if 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") - 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_reduction_input(self, input: torch.Tensor) -> None: - self._validate_tensor(input) - if input.dtype not in _REDUCTION_DTYPES: - raise TypeError( - "deterministic reductions support float32, float16, and bfloat16; " - f"got {input.dtype}" - ) - - def _validate_gather_input(self, input: torch.Tensor) -> None: - self._validate_tensor(input) - if input.dim() == 0: - raise ValueError("all_gather input must have at least one dimension") - - @staticmethod - def _validate_output( - output: torch.Tensor, - input: torch.Tensor, - output_shape: tuple[int, ...], - ) -> None: - if not isinstance(output, torch.Tensor): - raise TypeError(f"out must be a torch.Tensor, got {type(output)!r}") - 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 tuple(output.shape) != output_shape: - raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") - if not output.is_contiguous(): - raise ValueError("out must be contiguous") - - def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: - if self.world_size == 1: - return - signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) - if signature in self._validated_signatures: - return - 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}" - ) - self._validated_signatures.add(signature) - - def _validate_matching_capacity(self) -> None: - if self.world_size == 1: - return - capacities: list[int | None] = [None] * self.world_size - dist.all_gather_object(capacities, self.max_size_bytes, group=self.group) - if any(peer_capacity != self.max_size_bytes for peer_capacity in capacities): - raise ValueError(f"all ranks must use the same max_size_bytes; got {capacities}") - - def _all_gather_transport( - self, - input: torch.Tensor, - *, - gathered_flat: torch.Tensor | None = None, - ) -> torch.Tensor: - # Flattening makes the output contract independent of whether a given - # ProcessGroup implements the concatenation or stacking form of AG. - if self.world_size == 1: - if gathered_flat is None: - gathered_flat = input.clone().view(-1) - else: - gathered_flat.copy_(input.view(-1)) - return gathered_flat.reshape((1, *input.shape)) - input_flat = input.view(-1) - required_elements = self.world_size * input_flat.numel() - if gathered_flat is None: - gathered_flat = self._workspace_for(input, required_elements) - elif ( - gathered_flat.numel() != required_elements - or gathered_flat.dtype != input.dtype - or gathered_flat.device != input.device - or not gathered_flat.is_contiguous() - ): - raise ValueError("gathered transport output has an invalid layout") - if "nccl" in self._backend: - # PyTorch exposes RCCL through the NCCL ProcessGroup API. Keep - # this as a tensor-only transport; reduction happens below. - dist.all_gather_into_tensor(gathered_flat, input_flat, group=self.group) - else: - # Some reference backends (notably Gloo versions without - # all_gather_into_tensor) only implement the list API. - gathered_chunks = list( - gathered_flat.reshape(self.world_size, input_flat.numel()).unbind(0) - ) - dist.all_gather(gathered_chunks, input_flat, group=self.group) - return gathered_flat.reshape((self.world_size, *input.shape)) - - def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.Tensor: - required_bytes = required_elements * input.element_size() - workspace = self._workspace - if workspace is None or workspace.numel() < required_bytes: - workspace = torch.empty(required_bytes, dtype=torch.uint8, device=self.device) - self._workspace = workspace - # Tensor.view(dtype) reinterprets the aligned byte allocation without - # an allocation or copy. Restrict the view to the current operation. - return workspace[:required_bytes].view(input.dtype) - - @staticmethod - def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: - world_size = rank_inputs.size(0) - if world_size not in _SUPPORTED_WORLD_SIZES: - raise ValueError( - "balanced reduction requires rank inputs for world_size in " - f"{_SUPPORTED_WORLD_SIZES}, got {world_size}" - ) - # ``rank_inputs`` is the private transport workspace for reductions, - # so fixed-tree nodes can be accumulated in place. This preserves the - # exact pairings while avoiding one temporary allocation per tree node. - stride = 1 - while stride < world_size: - for index in range(0, world_size, 2 * stride): - rank_inputs[index].add_(rank_inputs[index + stride]) - stride *= 2 - return rank_inputs[0] - - -class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): - """ROCm collective using RCCL AllGather strictly as tensor transport.""" - - backend_id = "rccl_all_gather_balanced_tree" - - 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 getattr(torch.version, "hip", None) is None: - raise RuntimeError("RCCL deterministic collectives require a ROCm PyTorch build") - if not torch.cuda.is_available(): - raise RuntimeError("RCCL deterministic collectives require an available ROCm device") - if device is not None and not isinstance(device, int): - requested_device = torch.device(device) - if requested_device.type != "cuda": - raise ValueError( - f"RCCL deterministic collectives require a ROCm device, got {device!r}" - ) - super().__init__(group=group, device=device, max_size_bytes=max_size_bytes) - if self.device.type != "cuda": - raise ValueError( - f"RCCL deterministic collectives require a ROCm device, got {device!r}" - ) - if "nccl" not in self._backend: - raise RuntimeError( - "RCCL deterministic collectives require PyTorch's NCCL process-group API" - ) - - -def create_deterministic_collective( - group: dist.ProcessGroup | None = None, - device: torch.device | str | int | None = None, - *, - max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, -) -> Any: - """Create the platform-appropriate deterministic collective. - - CUDA keeps using the existing IPC implementation. ROCm uses RCCL only to - gather rank inputs, followed by the shared local balanced-tree reduction. - The returned object has independent ownership; callers may cache it by - process-group/rank/device and must close an entry before replacing it. - """ - - if getattr(torch.version, "hip", None) is not None: - return RCCLDeterministicCollective( - group=group, - device=device, - max_size_bytes=max_size_bytes, - ) - - # Import lazily to keep the existing CUDA implementation and its extension - # checks independent of the generic transport reference above. - from rl_engine.distributed.collectives import DeterministicCollective - - return DeterministicCollective( - group=group, - device=device, - max_size_bytes=max_size_bytes, - ) - - -__all__ = [ - "RCCLDeterministicCollective", - "TorchDistributedDeterministicCollective", - "create_deterministic_collective", -] diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py index ed1da194..dddd674b 100644 --- a/tests/distributed/test_transport_deterministic_collective.py +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -8,8 +8,8 @@ import pytest import torch -import rl_engine.distributed.collectives as cuda_collectives -import rl_engine.distributed.transport_collectives as transport_collectives +import rl_engine.distributed as distributed +import rl_engine.distributed.collectives as collectives from rl_engine.distributed import ( RCCLDeterministicCollective, TorchDistributedDeterministicCollective, @@ -92,6 +92,18 @@ def all_gather( destination.copy_(peer.reshape(-1)) +def test_public_exports_use_canonical_collectives_module() -> None: + assert distributed.DeterministicCollective is collectives.DeterministicCollective + assert distributed.RCCLDeterministicCollective is collectives.RCCLDeterministicCollective + assert ( + distributed.TorchDistributedDeterministicCollective + is collectives.TorchDistributedDeterministicCollective + ) + assert ( + distributed.create_deterministic_collective is collectives.create_deterministic_collective + ) + + def _make_collective( monkeypatch: pytest.MonkeyPatch, peer_inputs: list[torch.Tensor], @@ -109,7 +121,7 @@ def _make_collective( peer_signatures=peer_signatures, peer_capacities=peer_capacities, ) - monkeypatch.setattr(transport_collectives, "dist", fake_dist) + monkeypatch.setattr(collectives, "dist", fake_dist) collective = TorchDistributedDeterministicCollective( group=object(), device="cpu", @@ -312,14 +324,14 @@ def test_unsupported_world_size_is_rejected( world_size: int, ) -> None: fake_dist = _FakeDistributed([torch.ones(1)] * world_size) - monkeypatch.setattr(transport_collectives, "dist", fake_dist) + monkeypatch.setattr(collectives, "dist", fake_dist) with pytest.raises(ValueError, match="world_size in"): TorchDistributedDeterministicCollective(group=object(), device="cpu") def test_rccl_class_requires_rocm_build(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(transport_collectives.torch.version, "hip", None, raising=False) + monkeypatch.setattr(collectives.torch.version, "hip", None, raising=False) with pytest.raises(RuntimeError, match="ROCm PyTorch build"): RCCLDeterministicCollective(group=object(), device="cuda:0") @@ -327,10 +339,10 @@ def test_rccl_class_requires_rocm_build(monkeypatch: pytest.MonkeyPatch) -> None def test_rccl_class_requires_nccl_process_group(monkeypatch: pytest.MonkeyPatch) -> None: fake_dist = _FakeDistributed([torch.ones(1)], backend="gloo") - monkeypatch.setattr(transport_collectives, "dist", fake_dist) - monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) - monkeypatch.setattr(transport_collectives.torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(transport_collectives.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(collectives, "dist", fake_dist) + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(collectives.torch.cuda, "current_device", lambda: 0) with pytest.raises(RuntimeError, match="NCCL process-group API"): RCCLDeterministicCollective(group=object(), device="cuda:0") @@ -339,8 +351,8 @@ def test_rccl_class_requires_nccl_process_group(monkeypatch: pytest.MonkeyPatch) def test_rccl_class_rejects_cpu_before_process_group_exchange( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) - monkeypatch.setattr(transport_collectives.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives.torch.cuda, "is_available", lambda: True) with pytest.raises(ValueError, match="ROCm device"): RCCLDeterministicCollective(group=object(), device="cpu") @@ -354,11 +366,11 @@ def fake_rccl(**kwargs: Any) -> object: calls.append(kwargs) return sentinel - monkeypatch.setattr(transport_collectives.torch.version, "hip", "6.3", raising=False) - monkeypatch.setattr(transport_collectives, "RCCLDeterministicCollective", fake_rccl) + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives, "RCCLDeterministicCollective", fake_rccl) group = object() - result = transport_collectives.create_deterministic_collective( + result = collectives.create_deterministic_collective( group=group, device="cuda:3", max_size_bytes=1234, @@ -376,11 +388,11 @@ def fake_cuda(**kwargs: Any) -> object: calls.append(kwargs) return sentinel - monkeypatch.setattr(transport_collectives.torch.version, "hip", None, raising=False) - monkeypatch.setattr(cuda_collectives, "DeterministicCollective", fake_cuda) + monkeypatch.setattr(collectives.torch.version, "hip", None, raising=False) + monkeypatch.setattr(collectives, "DeterministicCollective", fake_cuda) group = object() - result = transport_collectives.create_deterministic_collective( + result = collectives.create_deterministic_collective( group=group, device="cuda:1", max_size_bytes=4321, From f0513d2c4769d3f757f6ff980864e39355f73c3a Mon Sep 17 00:00:00 2001 From: vensen Date: Sat, 29 Aug 2026 09:43:49 +0000 Subject: [PATCH 4/5] chore: trigger dco app Signed-off-by: vensen From 285dda1600a2bd09ce2ddd6d1df0c12437b5efe4 Mon Sep 17 00:00:00 2001 From: vensen Date: Sun, 30 Aug 2026 14:02:44 +0000 Subject: [PATCH 5/5] style(distributed): satisfy collective formatting Signed-off-by: vensen --- rl_engine/distributed/collectives.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 0d040ea3..119b22bb 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -630,8 +630,7 @@ def reduce_scatter_many( if not inputs: raise ValueError("reduce_scatter_many requires at least one input") return tuple( - self.reduce_scatter(input, validate_signature=validate_signature) - for input in inputs + self.reduce_scatter(input, validate_signature=validate_signature) for input in inputs ) def close(self) -> None: