diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 630dba3e..40226ca3 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -114,6 +114,54 @@ void deterministic_collective_all_gather_fused( int64_t handle, torch::Tensor& input, torch::Tensor& output); #endif +#if defined(KERNEL_ALIGN_WITH_ROCM) +// ROCm keeps arithmetic in a fixed balanced tree while using either RCCL or +// HIP IPC for rank-ordered transport. These kernels expose the local and IPC +// reduction paths without changing the CUDA implementation. +void deterministic_collective_rocm_all_reduce( + torch::Tensor rank_inputs, + torch::Tensor output); +void deterministic_collective_rocm_reduce_scatter( + torch::Tensor rank_inputs, + torch::Tensor output); +torch::Tensor deterministic_collective_rocm_ipc_allocate(int64_t size_bytes); +std::tuple, int64_t> +deterministic_collective_rocm_ipc_meta(torch::Tensor tensor); +int64_t deterministic_collective_rocm_ipc_create( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank); +void deterministic_collective_rocm_ipc_synchronize(int64_t handle); +void deterministic_collective_rocm_ipc_destroy(int64_t handle); +void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input); +void deterministic_collective_rocm_ipc_all_reduce( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_all_reduce_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter_many( + int64_t handle, + const std::vector& inputs, + const std::vector& outputs); +void deterministic_collective_rocm_ipc_all_gather( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_all_gather_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +#endif + // Batch-Invariant Deterministic GEMM Declarations bool det_gemm_sm90_compiled(); torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); @@ -469,6 +517,56 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Run a fused small-message deterministic rank-ordered all-gather"); #endif +#if defined(KERNEL_ALIGN_WITH_ROCM) + m.def( + "deterministic_collective_rocm_all_reduce", + &deterministic_collective_rocm_all_reduce, + "Run the ROCm fixed-tree all-reduce kernel"); + m.def( + "deterministic_collective_rocm_reduce_scatter", + &deterministic_collective_rocm_reduce_scatter, + "Run the ROCm fixed-tree reduce-scatter kernel"); + m.def("deterministic_collective_rocm_ipc_meta", + &deterministic_collective_rocm_ipc_meta, + "Export a ROCm allocation for IPC deterministic collectives"); + m.def("deterministic_collective_rocm_ipc_allocate", + &deterministic_collective_rocm_ipc_allocate, + "Allocate ROCm memory that supports IPC export"); + m.def("deterministic_collective_rocm_ipc_create", + &deterministic_collective_rocm_ipc_create, + "Create a ROCm IPC deterministic collective state"); + m.def("deterministic_collective_rocm_ipc_destroy", + &deterministic_collective_rocm_ipc_destroy, + "Destroy a ROCm IPC deterministic collective state"); + m.def("deterministic_collective_rocm_ipc_synchronize", + &deterministic_collective_rocm_ipc_synchronize, + "Wait until every rank finishes reading ROCm IPC staging"); + m.def("deterministic_collective_rocm_ipc_stage", + &deterministic_collective_rocm_ipc_stage, + "Stage an input for ROCm IPC deterministic collectives"); + m.def("deterministic_collective_rocm_ipc_all_reduce", + &deterministic_collective_rocm_ipc_all_reduce, + "Run a direct ROCm IPC fixed-tree all-reduce"); + m.def("deterministic_collective_rocm_ipc_all_reduce_input", + &deterministic_collective_rocm_ipc_all_reduce_input, + "Stage and run a direct ROCm IPC fixed-tree all-reduce"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter", + &deterministic_collective_rocm_ipc_reduce_scatter, + "Run a direct ROCm IPC fixed-tree reduce-scatter"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter_input", + &deterministic_collective_rocm_ipc_reduce_scatter_input, + "Stage and run a direct ROCm IPC fixed-tree reduce-scatter"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter_many", + &deterministic_collective_rocm_ipc_reduce_scatter_many, + "Run multiple ROCm IPC fixed-tree reduce-scatters with one synchronization"); + m.def("deterministic_collective_rocm_ipc_all_gather", + &deterministic_collective_rocm_ipc_all_gather, + "Run a direct ROCm IPC rank-ordered all-gather"); + m.def("deterministic_collective_rocm_ipc_all_gather_input", + &deterministic_collective_rocm_ipc_all_gather_input, + "Stage and run a direct ROCm IPC rank-ordered all-gather"); +#endif + // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. #if !defined(USE_ROCM) m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); diff --git a/csrc/rocm/distributed/deterministic_collective.hip b/csrc/rocm/distributed/deterministic_collective.hip new file mode 100644 index 00000000..c8448300 --- /dev/null +++ b/csrc/rocm/distributed/deterministic_collective.hip @@ -0,0 +1,1099 @@ +// ROCm fixed-tree reduction kernels for the RCCL transport collective. +// +// RCCL is intentionally used only to transport rank-ordered tensors. The +// kernels below perform the arithmetic locally in the exact balanced tree used +// by the Python reference implementation. ReduceScatter receives a view that +// contains only the destination rank's shard, so it does not reduce unrelated +// rows. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 4096; +constexpr int kMaxWorldSize = 8; +constexpr int64_t kIPCControlBytes = 256; +constexpr int64_t kIPCReadyOffset = 0; +constexpr int64_t kIPCDoneOffset = 64; +constexpr int64_t kIPCCloseOffset = 128; + +struct PeerPointers { + const void* values[kMaxWorldSize]; +}; + +struct PeerSignals { + uint64_t* ready[kMaxWorldSize]; + uint64_t* done[kMaxWorldSize]; + uint64_t* closed[kMaxWorldSize]; +}; + +template +__device__ __forceinline__ scalar_t ordered_add(scalar_t lower, scalar_t upper) { + // Keep every parent as a separate expression. ROCm builds do not enable + // fast-math, so this is the same dtype operation as torch.add_ for the + // supported floating-point dtypes. + return lower + upper; +} + +template +__device__ __forceinline__ scalar_t fixed_tree_reduce( + const scalar_t* values, + int64_t rank_stride, + int64_t index) { + static_assert( + WorldSize == 1 || WorldSize == 2 || WorldSize == 4 || WorldSize == 8, + "unsupported deterministic collective world size"); + if constexpr (WorldSize == 1) { + return values[index]; + } else { + const scalar_t sum01 = ordered_add( + values[index], + values[rank_stride + index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const scalar_t sum23 = ordered_add( + values[2 * rank_stride + index], + values[3 * rank_stride + index]); + const scalar_t sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const scalar_t sum45 = ordered_add( + values[4 * rank_stride + index], + values[5 * rank_stride + index]); + const scalar_t sum67 = ordered_add( + values[6 * rank_stride + index], + values[7 * rank_stride + index]); + const scalar_t sum47 = ordered_add(sum45, sum67); + return ordered_add(sum03, sum47); + } + } + } +} + +template +__global__ void fixed_tree_reduce_kernel( + const scalar_t* __restrict__ values, + scalar_t* __restrict__ output, + int64_t rank_stride, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = fixed_tree_reduce(values, rank_stride, index); + } +} + +template +void launch_fixed_tree_reduce( + const scalar_t* values, + scalar_t* output, + int64_t rank_stride, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 2: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 4: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 8: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void validate_inputs( + const torch::Tensor& rank_inputs, + const torch::Tensor& output, + const char* name) { + TORCH_CHECK(rank_inputs.is_cuda(), name, ": rank_inputs must be a CUDA/ROCm tensor"); + TORCH_CHECK(output.is_cuda(), name, ": output must be a CUDA/ROCm tensor"); + TORCH_CHECK(rank_inputs.scalar_type() == output.scalar_type(), name, ": dtype mismatch"); + TORCH_CHECK(rank_inputs.dim() >= 1, name, ": rank_inputs must have a rank dimension"); + TORCH_CHECK(rank_inputs.size(0) == 1 || rank_inputs.size(0) == 2 || + rank_inputs.size(0) == 4 || rank_inputs.size(0) == 8, + name, ": unsupported rank dimension ", rank_inputs.size(0)); + TORCH_CHECK(rank_inputs.select(0, 0).is_contiguous(), + name, ": each rank slice must be contiguous"); + TORCH_CHECK(output.is_contiguous(), name, ": output must be contiguous"); + TORCH_CHECK(rank_inputs.device() == output.device(), name, ": device mismatch"); + TORCH_CHECK(rank_inputs.numel() == output.numel() * rank_inputs.size(0), + name, ": rank_inputs/output element count mismatch"); +} + +void launch_dispatch( + const torch::Tensor& rank_inputs, + const torch::Tensor& output, + const char* name) { + validate_inputs(rank_inputs, output, name); + const int64_t world_size = rank_inputs.size(0); + const int64_t element_count = output.numel(); + if (element_count == 0) { + return; + } + const int64_t rank_stride = rank_inputs.stride(0); + const auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + rank_inputs.scalar_type(), + "deterministic_collective_rocm_fixed_tree", + [&] { + launch_fixed_tree_reduce( + rank_inputs.data_ptr(), + output.data_ptr(), + rank_stride, + element_count, + world_size, + stream); + }); +} + +template +__device__ __forceinline__ scalar_t ipc_fixed_tree_reduce( + const PeerPointers& peers, + int64_t index) { + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const scalar_t sum01 = ordered_add(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const scalar_t sum23 = ordered_add(rank2[index], rank3[index]); + const scalar_t sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const scalar_t sum45 = ordered_add(rank4[index], rank5[index]); + const scalar_t sum67 = ordered_add(rank6[index], rank7[index]); + return ordered_add(sum03, ordered_add(sum45, sum67)); + } + } + } +} + +template +__device__ __forceinline__ packed_t ordered_add_packed( + packed_t lower, + packed_t upper); + +template <> +__device__ __forceinline__ __half2 ordered_add_packed( + __half2 lower, + __half2 upper) { + return __hadd2(lower, upper); +} + +template <> +__device__ __forceinline__ __hip_bfloat162 ordered_add_packed( + __hip_bfloat162 lower, + __hip_bfloat162 upper) { + return __hadd2(lower, upper); +} + +template +__device__ __forceinline__ packed_t ipc_fixed_tree_reduce_packed( + const PeerPointers& peers, + int64_t index) { + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const packed_t sum01 = ordered_add_packed(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const packed_t sum23 = ordered_add_packed(rank2[index], rank3[index]); + const packed_t sum03 = ordered_add_packed(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const packed_t sum45 = ordered_add_packed(rank4[index], rank5[index]); + const packed_t sum67 = ordered_add_packed(rank6[index], rank7[index]); + return ordered_add_packed(sum03, ordered_add_packed(sum45, sum67)); + } + } + } +} + +template +__global__ void ipc_fixed_tree_reduce_kernel( + PeerPointers peers, + scalar_t* __restrict__ output, + int64_t input_offset, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = ipc_fixed_tree_reduce( + peers, + input_offset + index); + } +} + +template +__global__ void ipc_fixed_tree_reduce_packed_kernel( + PeerPointers peers, + packed_t* __restrict__ output, + int64_t input_offset, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = ipc_fixed_tree_reduce_packed( + peers, + input_offset + index); + } +} + +template +void launch_ipc_fixed_tree_reduce( + const PeerPointers& peers, + scalar_t* output, + int64_t input_offset, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 2: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 4: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 8: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_ipc_fixed_tree_reduce_packed( + const PeerPointers& peers, + packed_t* output, + int64_t input_offset, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 2: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 4: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 8: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +__global__ void ipc_wait_signal_kernel( + PeerSignals signals, + uint64_t sequence, + int64_t world_size, + bool wait_for_done) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + for (int peer = 0; peer < world_size; ++peer) { + uint64_t* signal = wait_for_done ? signals.done[peer] : signals.ready[peer]; + while (__hip_atomic_load( + signal, + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + +__global__ void ipc_mark_signal_kernel(uint64_t* signal, uint64_t sequence) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + __hip_atomic_store( + signal, + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +__global__ void ipc_mark_ready_and_wait_kernel( + PeerSignals signals, + int64_t rank, + uint64_t sequence, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + __hip_atomic_store( + signals.ready[rank], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.ready[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + +__global__ void ipc_close_and_wait_kernel( + PeerSignals signals, + int64_t rank, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + for (int peer = 0; peer < world_size; ++peer) { + __hip_atomic_fetch_add( + signals.closed[peer], + static_cast(1), + __ATOMIC_ACQ_REL, + __HIP_MEMORY_SCOPE_SYSTEM); + } + while (__hip_atomic_load( + signals.closed[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < static_cast(world_size)) { + __builtin_amdgcn_s_sleep(1); + } +} + +__global__ void ipc_all_gather_uint4_kernel( + PeerPointers peers, + uint4* __restrict__ output, + int64_t vectors_per_rank, + int64_t world_size) { + const int64_t total_vectors = vectors_per_rank * world_size; + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < total_vectors; index += stride) { + const int peer = static_cast(index / vectors_per_rank); + const int64_t peer_index = index - static_cast(peer) * vectors_per_rank; + output[index] = static_cast(peers.values[peer])[peer_index]; + } +} + +__global__ void ipc_all_gather_bytes_kernel( + PeerPointers peers, + uint8_t* __restrict__ output, + int64_t bytes_per_rank, + int64_t world_size) { + const int64_t total_bytes = bytes_per_rank * world_size; + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < total_bytes; index += stride) { + const int peer = static_cast(index / bytes_per_rank); + const int64_t peer_index = index - static_cast(peer) * bytes_per_rank; + output[index] = static_cast(peers.values[peer])[peer_index]; + } +} + +void launch_ipc_all_gather( + const PeerPointers& peers, + void* output, + int64_t bytes_per_rank, + int64_t world_size, + hipStream_t stream) { + if (bytes_per_rank == 0) { + return; + } + if (bytes_per_rank % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(output) % alignof(uint4) == 0) { + const int64_t vectors_per_rank = bytes_per_rank / sizeof(uint4); + const int64_t total_vectors = vectors_per_rank * world_size; + const int blocks = static_cast(std::min( + kMaxBlocks, + (total_vectors + kThreads - 1) / kThreads)); + hipLaunchKernelGGL( + ipc_all_gather_uint4_kernel, + dim3(blocks), + dim3(kThreads), + 0, + stream, + peers, + static_cast(output), + vectors_per_rank, + world_size); + } else { + const int64_t total_bytes = bytes_per_rank * world_size; + const int blocks = static_cast(std::min( + kMaxBlocks, + (total_bytes + kThreads - 1) / kThreads)); + hipLaunchKernelGGL( + ipc_all_gather_bytes_kernel, + dim3(blocks), + dim3(kThreads), + 0, + stream, + peers, + static_cast(output), + bytes_per_rank, + world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_wait_signal( + const PeerSignals& signals, + uint64_t sequence, + int64_t world_size, + bool wait_for_done, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_wait_signal_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + sequence, + world_size, + wait_for_done); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_signal(uint64_t* signal, uint64_t sequence, hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_signal_kernel, + dim3(1), + dim3(1), + 0, + stream, + signal, + sequence); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_ready_and_wait( + const PeerSignals& signals, + int64_t rank, + uint64_t sequence, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_ready_and_wait_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + sequence, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_close_and_wait( + const PeerSignals& signals, + int64_t rank, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_close_and_wait_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +class ROCmIPCCollectiveState { + public: + ROCmIPCCollectiveState( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) + : rank_(rank), + world_size_(static_cast(handles.size())), + device_index_(staging.get_device()), + capacity_bytes_(staging.numel() * staging.element_size() - kIPCControlBytes) { + TORCH_CHECK(staging.is_cuda(), "ROCm IPC staging buffer must be on device"); + TORCH_CHECK(staging.is_contiguous(), "ROCm IPC staging buffer must be contiguous"); + TORCH_CHECK(staging.scalar_type() == torch::kUInt8, + "ROCm IPC staging buffer must have dtype uint8"); + TORCH_CHECK(capacity_bytes_ > 0, "ROCm IPC staging capacity must be positive"); + TORCH_CHECK( + world_size_ == 1 || world_size_ == 2 || world_size_ == 4 || world_size_ == 8, + "ROCm IPC deterministic collectives require world size 1, 2, 4, or 8"); + TORCH_CHECK(offsets.size() == handles.size(), "one IPC offset is required per rank"); + TORCH_CHECK(rank_ >= 0 && rank_ < world_size_, "invalid ROCm IPC rank"); + + set_peer_pointers(rank_, staging.data_ptr()); + try { + for (int peer = 0; peer < world_size_; ++peer) { + if (peer == rank_) { + continue; + } + TORCH_CHECK(handles[peer].size() == sizeof(hipIpcMemHandle_t), + "invalid ROCm IPC handle size for rank ", peer); + TORCH_CHECK(offsets[peer] >= 0, "invalid negative ROCm IPC offset"); + hipIpcMemHandle_t handle{}; + auto* raw_handle = reinterpret_cast(&handle); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + TORCH_CHECK(handles[peer][byte] >= 0 && handles[peer][byte] <= 255, + "invalid ROCm IPC handle byte for rank ", peer); + raw_handle[byte] = static_cast(handles[peer][byte]); + } + void* base = nullptr; + C10_HIP_CHECK(hipIpcOpenMemHandle( + &base, + handle, + hipIpcMemLazyEnablePeerAccess)); + imported_bases_[peer] = base; + set_peer_pointers( + peer, + static_cast(base) + offsets[peer]); + } + } catch (...) { + close_imports(); + throw; + } + } + + ~ROCmIPCCollectiveState() { + int previous_device = -1; + if (hipGetDevice(&previous_device) == hipSuccess && previous_device != device_index_) { + if (hipSetDevice(device_index_) != hipSuccess) { + return; + } + } + close_imports(); + if (previous_device >= 0 && previous_device != device_index_) { + C10_CUDA_IGNORE_ERROR(hipSetDevice(previous_device)); + } + } + + int device_index() const { + return device_index_; + } + + void stage(torch::Tensor input, hipStream_t stream) { + check_tensor(input, "input"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK(input_bytes <= capacity_bytes_, + "input exceeds ROCm IPC staging capacity"); + ++sequence_; + launch_wait_signal( + signals_, + sequence_ - 1, + world_size_, + true, + stream); + if (input_bytes > 0) { + C10_HIP_CHECK(hipMemcpyAsync( + const_cast(peers_.values[rank_]), + input.data_ptr(), + input_bytes, + hipMemcpyDeviceToDevice, + stream)); + } + launch_mark_ready_and_wait( + signals_, + rank_, + sequence_, + world_size_, + stream); + staged_bytes_ = input_bytes; + staged_type_ = input.scalar_type(); + } + + void all_reduce(torch::Tensor output, hipStream_t stream) const { + check_reduction_output(output, staged_bytes_, "all_reduce"); + launch(output, 0, output.numel(), stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void reduce_scatter(torch::Tensor output, hipStream_t stream) const { + check_reduction_output( + output, + staged_bytes_ / world_size_, + "reduce_scatter"); + launch( + output, + rank_ * output.numel(), + output.numel(), + stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void reduce_scatter_many( + const std::vector& inputs, + const std::vector& outputs, + hipStream_t stream) { + TORCH_CHECK(!inputs.empty(), "reduce_scatter_many requires at least one input"); + TORCH_CHECK(inputs.size() == outputs.size(), + "reduce_scatter_many input/output count mismatch"); + + int64_t total_bytes = 0; + const auto scalar_type = inputs.front().scalar_type(); + for (size_t index = 0; index < inputs.size(); ++index) { + const auto& input = inputs[index]; + const auto& output = outputs[index]; + check_tensor(input, "reduce_scatter_many input"); + check_tensor(output, "reduce_scatter_many output"); + TORCH_CHECK(input.scalar_type() == scalar_type && output.scalar_type() == scalar_type, + "reduce_scatter_many dtype mismatch"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK(input_bytes == output.numel() * output.element_size() * world_size_, + "reduce_scatter_many output size mismatch"); + TORCH_CHECK(input_bytes <= capacity_bytes_ - total_bytes, + "reduce_scatter_many inputs exceed ROCm IPC staging capacity"); + total_bytes += input_bytes; + } + + ++sequence_; + launch_wait_signal(signals_, sequence_ - 1, world_size_, true, stream); + int64_t byte_offset = 0; + for (const auto& input : inputs) { + const int64_t input_bytes = input.numel() * input.element_size(); + if (input_bytes > 0) { + C10_HIP_CHECK(hipMemcpyAsync( + static_cast(const_cast(peers_.values[rank_])) + byte_offset, + input.data_ptr(), + input_bytes, + hipMemcpyDeviceToDevice, + stream)); + } + byte_offset += input_bytes; + } + launch_mark_ready_and_wait( + signals_, + rank_, + sequence_, + world_size_, + stream); + + int64_t element_offset = 0; + for (size_t index = 0; index < outputs.size(); ++index) { + const auto& input = inputs[index]; + const auto& output = outputs[index]; + launch( + output, + element_offset + rank_ * output.numel(), + output.numel(), + stream); + element_offset += input.numel(); + } + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void all_gather(torch::Tensor output, hipStream_t stream) const { + check_tensor(output, "all_gather"); + TORCH_CHECK(staged_type_ != at::ScalarType::Undefined, "stage must be called first"); + TORCH_CHECK(output.scalar_type() == staged_type_, "all_gather dtype mismatch"); + TORCH_CHECK( + output.numel() * output.element_size() == staged_bytes_ * world_size_, + "all_gather output size mismatch"); + launch_ipc_all_gather( + peers_, + output.data_ptr(), + staged_bytes_, + world_size_, + stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void synchronize(hipStream_t stream) const { + launch_wait_signal(signals_, sequence_, world_size_, true, stream); + launch_close_and_wait(signals_, rank_, world_size_, stream); + } + + private: + void set_peer_pointers(int peer, void* allocation_base) { + auto* bytes = static_cast(allocation_base); + signals_.ready[peer] = reinterpret_cast(bytes + kIPCReadyOffset); + signals_.done[peer] = reinterpret_cast(bytes + kIPCDoneOffset); + signals_.closed[peer] = reinterpret_cast(bytes + kIPCCloseOffset); + peers_.values[peer] = bytes + kIPCControlBytes; + } + + void check_tensor(const torch::Tensor& tensor, const char* name) const { + TORCH_CHECK(tensor.is_cuda(), name, " must be a ROCm tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.get_device() == device_index_, name, " device mismatch"); + } + + void check_reduction_output( + const torch::Tensor& output, + int64_t expected_bytes, + const char* name) const { + check_tensor(output, name); + TORCH_CHECK(staged_type_ != at::ScalarType::Undefined, "stage must be called first"); + TORCH_CHECK(output.scalar_type() == staged_type_, name, " dtype mismatch"); + TORCH_CHECK(output.numel() * output.element_size() == expected_bytes, + name, " output size mismatch"); + } + + void launch( + torch::Tensor output, + int64_t input_offset, + int64_t element_count, + hipStream_t stream) const { + if (element_count == 0) { + return; + } + if (element_count % 2 == 0 && input_offset % 2 == 0) { + if (output.scalar_type() == at::ScalarType::Half && + reinterpret_cast(output.data_ptr()) % alignof(__half2) == 0) { + launch_ipc_fixed_tree_reduce_packed<__half2>( + peers_, + reinterpret_cast<__half2*>(output.data_ptr()), + input_offset / 2, + element_count / 2, + world_size_, + stream); + return; + } + if (output.scalar_type() == at::ScalarType::BFloat16 && + reinterpret_cast(output.data_ptr()) % + alignof(__hip_bfloat162) == + 0) { + launch_ipc_fixed_tree_reduce_packed<__hip_bfloat162>( + peers_, + reinterpret_cast<__hip_bfloat162*>(output.data_ptr()), + input_offset / 2, + element_count / 2, + world_size_, + stream); + return; + } + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + output.scalar_type(), + "deterministic_collective_rocm_ipc_fixed_tree", + [&] { + launch_ipc_fixed_tree_reduce( + peers_, + output.data_ptr(), + input_offset, + element_count, + world_size_, + stream); + }); + } + + void close_imports() noexcept { + for (int peer = 0; peer < world_size_; ++peer) { + if (imported_bases_[peer] != nullptr) { + C10_CUDA_IGNORE_ERROR(hipIpcCloseMemHandle(imported_bases_[peer])); + imported_bases_[peer] = nullptr; + } + } + } + + int64_t rank_; + int64_t world_size_; + int device_index_; + int64_t capacity_bytes_; + int64_t staged_bytes_{0}; + at::ScalarType staged_type_{at::ScalarType::Undefined}; + uint64_t sequence_{0}; + PeerPointers peers_{}; + PeerSignals signals_{}; + std::array imported_bases_{}; +}; + +ROCmIPCCollectiveState* ipc_state(int64_t handle) { + TORCH_CHECK(handle != 0, "ROCm IPC collective handle is closed"); + return reinterpret_cast(handle); +} + +} // namespace + +void deterministic_collective_rocm_all_reduce( + torch::Tensor rank_inputs, + torch::Tensor output) { + launch_dispatch(rank_inputs, output, "deterministic_collective_rocm_all_reduce"); +} + +void deterministic_collective_rocm_reduce_scatter( + torch::Tensor rank_inputs, + torch::Tensor output) { + launch_dispatch(rank_inputs, output, "deterministic_collective_rocm_reduce_scatter"); +} + +torch::Tensor deterministic_collective_rocm_ipc_allocate(int64_t size_bytes) { + TORCH_CHECK(size_bytes > 0, "ROCm IPC allocation size must be positive"); + int device_index = -1; + C10_HIP_CHECK(hipGetDevice(&device_index)); + const int64_t allocation_bytes = size_bytes + kIPCControlBytes; + void* pointer = nullptr; + C10_HIP_CHECK(hipMalloc(&pointer, static_cast(allocation_bytes))); + C10_HIP_CHECK(hipMemset(pointer, 0, static_cast(kIPCControlBytes))); + const auto options = torch::TensorOptions() + .dtype(torch::kUInt8) + .device(torch::Device(torch::kCUDA, device_index)); + return torch::from_blob( + pointer, + {allocation_bytes}, + [device_index](void* allocation) { + int previous_device = -1; + if (hipGetDevice(&previous_device) != hipSuccess) { + return; + } + if (previous_device != device_index && hipSetDevice(device_index) != hipSuccess) { + return; + } + C10_CUDA_IGNORE_ERROR(hipFree(allocation)); + if (previous_device != device_index) { + C10_CUDA_IGNORE_ERROR(hipSetDevice(previous_device)); + } + }, + options); +} + +std::tuple, int64_t> +deterministic_collective_rocm_ipc_meta(torch::Tensor tensor) { + const c10::cuda::CUDAGuard device_guard(tensor.device()); + TORCH_CHECK(tensor.is_cuda(), "ROCm IPC tensor must be on device"); + TORCH_CHECK(tensor.is_contiguous(), "ROCm IPC tensor must be contiguous"); + TORCH_CHECK(tensor.numel() > 0, "ROCm IPC tensor must be non-empty"); + + hipIpcMemHandle_t handle{}; + const hipError_t export_error = hipIpcGetMemHandle(&handle, tensor.data_ptr()); + TORCH_CHECK( + export_error == hipSuccess, + "hipIpcGetMemHandle failed: ", + hipGetErrorString(export_error)); + const auto* raw_handle = reinterpret_cast(&handle); + std::vector bytes(sizeof(handle)); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + bytes[byte] = raw_handle[byte]; + } + return std::make_tuple(bytes, 0); +} + +int64_t deterministic_collective_rocm_ipc_create( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) { + const c10::cuda::CUDAGuard device_guard(staging.device()); + auto state = std::make_unique( + staging, + handles, + offsets, + rank); + return reinterpret_cast(state.release()); +} + +void deterministic_collective_rocm_ipc_destroy(int64_t handle) { + delete ipc_state(handle); +} + +void deterministic_collective_rocm_ipc_synchronize(int64_t handle) { + auto* state = ipc_state(handle); + const c10::cuda::CUDAGuard device_guard( + torch::Device(torch::kCUDA, state->device_index())); + const auto stream = at::cuda::getCurrentCUDAStream(); + state->synchronize(stream); +} + +void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->stage(input, stream); +} + +void deterministic_collective_rocm_ipc_all_reduce( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->all_reduce(output, stream); +} + +void deterministic_collective_rocm_ipc_all_reduce_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->all_reduce(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->reduce_scatter(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->reduce_scatter(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter_many( + int64_t handle, + const std::vector& inputs, + const std::vector& outputs) { + TORCH_CHECK(!inputs.empty(), "reduce_scatter_many requires at least one input"); + const c10::cuda::CUDAGuard device_guard(inputs.front().device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->reduce_scatter_many(inputs, outputs, stream); +} + +void deterministic_collective_rocm_ipc_all_gather( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->all_gather(output, stream); +} + +void deterministic_collective_rocm_ipc_all_gather_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->all_gather(output, stream); +} diff --git a/docs/design/rocm-deterministic-collectives.md b/docs/design/rocm-deterministic-collectives.md index ce104332..c105a537 100644 --- a/docs/design/rocm-deterministic-collectives.md +++ b/docs/design/rocm-deterministic-collectives.md @@ -12,12 +12,42 @@ 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. +Single-node ROCm uses a dedicated `hipMalloc` staging allocation on every rank. +The handles are exchanged once during construction and imported with HIP IPC. +Each call copies its local input into staging, publishes a system-scope GPU +sequence flag, and waits for every peer's matching sequence. The HIP kernel +then reads rank-ordered peer memory and evaluates exactly +`((rank0 + rank1) + (rank2 + rank3)) + ...`. A second sequence flag prevents a +rank from overwriting staging until every peer has finished reading it. + +FP16 and BF16 use two-element vector loads and `hadd2`. This changes only the +number of elements carried by an instruction: every scalar element retains the +same dtype, rank order, and expression grouping. FP32 and unaligned tails use +the scalar kernel. The executable Python tree remains the fallback for +CPU/reference backends and extensions built without the optional HIP source. + +The measured MI300X routing policy is: + +- AllReduce up to 768 KiB uses the direct IPC fixed-tree kernel. +- AllReduce from 768 KiB to 2.125 MiB uses the rank-major RCCL transport fallback. +- AllReduce at 2.125 MiB and above performs IPC ReduceScatter followed by RCCL + AllGather of the already-reduced shards. +- AllGather up to 256 KiB uses IPC peer copies; larger messages use RCCL. +- ReduceScatter uses IPC for all supported sizes and reduces only the local + destination shard. + +The ready publication and peer wait share one GPU kernel. The store is a +system-scope release and every peer load is a system-scope acquire. Done flags +remain a separate generation barrier because they protect staging reuse. At +close, every rank atomically acknowledges every peer allocation and waits for +all acknowledgements before releasing its local staging memory. + +Sequence-parallel FFN backward has two independent ReduceScatter lanes (gate +and up input gradients). On IPC, `reduce_scatter_many` copies the lanes into +disjoint staging ranges under one ready/done generation and launches one fixed +tree per output lane. It does not concatenate the inputs or mix their trees. +The RCCL fallback retains a measured packed-payload crossover because a larger +rank-major AllGather can be slower than two smaller calls. RCCL's `all_reduce` and `reduce_scatter` are not used for strict reductions. They guarantee a mathematical reduction but do not expose a stable @@ -51,7 +81,7 @@ provenance and toleranced correctness tests. ## Compute/communication fusion -The current ROCm collective is synchronous and reports +The current ROCm collective is stream ordered and reports `supports_async_overlap = False` and `supports_compute_communication_fusion = False`. FFN and Attention keep the dependency boundaries explicit: @@ -83,10 +113,9 @@ 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 +The IPC path favors a fixed arithmetic tree over native-RCCL reduction speed. +Large AllReduce avoids reducing the full tensor on every rank, but 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; @@ -95,8 +124,8 @@ world sizes 2/4/8 and representative FFN tensors: - 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. +Multi-node or unsupported IPC configurations fall back behind the same factory +after topology and symbol checks fail closed. Run the included native-RCCL comparison on a single node, for example: diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..b60169e1 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -22,6 +22,14 @@ def deterministic_collective_all_gather(handle: int, output: torch.Tensor) -> No def deterministic_collective_all_gather_fused( handle: int, input: torch.Tensor, output: torch.Tensor ) -> None: ... +def deterministic_collective_rocm_all_reduce( + rank_inputs: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_reduce_scatter( + rank_inputs: torch.Tensor, + output: torch.Tensor, +) -> None: ... def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... def batch_invariant_logp_sm90( @@ -222,3 +230,46 @@ def rmsnorm_backward_dw( rstd: torch.Tensor, mask: torch.Tensor, ) -> torch.Tensor: ... +def deterministic_collective_rocm_ipc_allocate(size_bytes: int) -> torch.Tensor: ... +def deterministic_collective_rocm_ipc_meta(tensor: torch.Tensor) -> tuple[list[int], int]: ... +def deterministic_collective_rocm_ipc_create( + staging: torch.Tensor, + handles: list[list[int]], + offsets: list[int], + rank: int, +) -> int: ... +def deterministic_collective_rocm_ipc_synchronize(handle: int) -> None: ... +def deterministic_collective_rocm_ipc_destroy(handle: int) -> None: ... +def deterministic_collective_rocm_ipc_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_rocm_ipc_all_reduce( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_all_reduce_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter_many( + handle: int, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], +) -> None: ... +def deterministic_collective_rocm_ipc_all_gather( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_all_gather_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 119b22bb..b599ad76 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic collectives for native CUDA IPC and rank-ordered transport. +"""Deterministic collectives for CUDA IPC and ROCm 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. +ROCm uses HIP IPC where it wins and RCCL otherwise. Reduction arithmetic stays +outside RCCL and follows the same fixed balanced rank tree on every rank. """ from __future__ import annotations @@ -19,6 +18,14 @@ _SUPPORTED_WORLD_SIZES = (1, 2, 4, 8) _DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +# Packing two independent lanes saves a collective launch for small tensors, +# but doubles the message size seen by RCCL. On MI300X, separate AllGather +# transports win once the packed payload reaches the multi-megabyte regime. +# Keep the crossover explicit and easy to retune with new RCCL releases. +_PACKED_REDUCE_SCATTER_MAX_BYTES = 8 * 1024 * 1024 +_ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES = 768 * 1024 +_ROCM_IPC_SHARDED_ALL_REDUCE_MIN_BYTES = 2176 * 1024 +_ROCM_IPC_ALL_GATHER_MAX_BYTES = 256 * 1024 _COLLECTIVE_STAGING_FRAMES = 3 _COLLECTIVE_FRAME_METADATA_BYTES = 3 * 8 _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) @@ -276,17 +283,33 @@ def reduce_scatter( def reduce_scatter_many( self, - inputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...] | list[torch.Tensor], *, + outs: tuple[torch.Tensor, ...] | list[torch.Tensor] | None = None, validate_signature: bool = True, ) -> tuple[torch.Tensor, ...]: - """Reduce-scatter several tensors through the single-tensor ABI.""" + """Compatibility fallback for CUDA IPC collectives. - if not inputs: + The native CUDA IPC backend has no packed transport primitive yet, so + it preserves its established behavior by issuing the individual + fixed-tree calls. The ROCm transport subclass overrides this method + with a packed implementation. + """ + + values = tuple(inputs) + if not values: raise ValueError("reduce_scatter_many requires at least one input") - return tuple( - self.reduce_scatter(input, validate_signature=validate_signature) for input in inputs + if outs is not None and len(outs) != len(values): + raise ValueError("reduce_scatter_many outs must match the number of inputs") + results = tuple( + self.reduce_scatter( + value, + out=None if outs is None else outs[index], + validate_signature=validate_signature, + ) + for index, value in enumerate(values) ) + return results def close(self) -> None: """Release imported CUDA IPC mappings after the last collective call.""" @@ -473,6 +496,10 @@ def __init__( self._backend = str(dist.get_backend(self.group)).lower() self._lock = threading.Lock() self._closed = False + # Keep a lifecycle marker for callers that historically inspected the + # CUDA IPC collective's ``_handle`` while managing the cache. Concrete + # transports own any native resource through their own state. + self._handle = id(self) # 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 @@ -538,14 +565,24 @@ def all_reduce( if out is None: out = torch.empty_like(input) self._validate_output(out, input, tuple(input.shape)) + if self.world_size == 1: + out.copy_(input) + return out with self._lock: self._check_open() if validate_signature: self._validate_matching_signature("all_reduce", input) + if self._direct_all_reduce(input, out): + return out rank_inputs = self._all_gather_transport(input) - reduced = self._balanced_tree_sum(rank_inputs) - out.copy_(reduced) + if not self._fused_reduction( + rank_inputs, + out, + operation="all_reduce", + ): + reduced = self._balanced_tree_sum(rank_inputs) + out.copy_(reduced) return out def all_gather( @@ -563,11 +600,16 @@ def all_gather( if out is None: out = torch.empty(output_shape, dtype=input.dtype, device=input.device) self._validate_output(out, input, output_shape) + if self.world_size == 1: + out.copy_(input) + return out with self._lock: self._check_open() if validate_signature: self._validate_matching_signature("all_gather", input) + if self._direct_all_gather(input, out): + return out self._all_gather_transport(input, gathered_flat=out.view(-1)) return out @@ -608,30 +650,182 @@ def reduce_scatter( if out is None: out = torch.empty(output_shape, dtype=input.dtype, device=input.device) self._validate_output(out, input, output_shape) + if self.world_size == 1: + out.copy_(input) + return out with self._lock: self._check_open() if validate_signature: self._validate_matching_signature("reduce_scatter", input) + if self._direct_reduce_scatter(input, out): + return out 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)) + # Only this rank's output shard participates in the reduction. The + # previous implementation reduced every global row and sliced the + # result afterwards, doing world_size times more arithmetic than + # ReduceScatter needs. The fixed rank tree is unchanged. + reduced = rank_inputs[:, begin : begin + rows_per_rank] + if not self._fused_reduction(reduced, out, operation="reduce_scatter"): + reduced = self._balanced_tree_sum(reduced) + out.copy_(reduced) return out def reduce_scatter_many( self, - inputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...] | list[torch.Tensor], *, + outs: tuple[torch.Tensor, ...] | list[torch.Tensor] | None = None, validate_signature: bool = True, ) -> tuple[torch.Tensor, ...]: - """Reduce-scatter several tensors through the single-tensor transport ABI.""" + """Reduce-scatter independent tensors in one fixed-tree collective. + + The tensors are packed along their final dimension, so each tensor's + element still follows the same balanced rank tree as an individual + ``reduce_scatter`` call. This is useful for independent gradient lanes: + packing them together removes one RCCL launch without changing the + floating-point expression for either lane. Inputs must have matching + shape/device/dtype except for the final dimension. + """ - if not inputs: + self._check_open() + values = tuple(inputs) + if not values: raise ValueError("reduce_scatter_many requires at least one input") - return tuple( - self.reduce_scatter(input, validate_signature=validate_signature) for input in inputs + if outs is not None and len(outs) != len(values): + raise ValueError("reduce_scatter_many outs must match the number of inputs") + if len(values) == 1: + return ( + self.reduce_scatter( + values[0], + out=None if outs is None else outs[0], + validate_signature=validate_signature, + ), + ) + + first = values[0] + self._validate_reduction_input(first) + if first.dim() < 2: + raise ValueError( + "reduce_scatter_many inputs must have at least two dimensions " + "when packing independent lanes" + ) + if first.size(0) % self.world_size != 0: + raise ValueError("reduce_scatter_many inputs must have a divisible leading dimension") + for value in values[1:]: + self._validate_reduction_input(value) + if value.dim() != first.dim() or value.shape[:-1] != first.shape[:-1]: + raise ValueError( + "reduce_scatter_many inputs must match in rank and all dimensions " + "except the final dimension" + ) + if value.device != first.device or value.dtype != first.dtype: + raise ValueError("reduce_scatter_many inputs must share device and dtype") + lane_sizes = tuple(int(value.size(-1)) for value in values) + rows_per_rank = first.size(0) // self.world_size + output_shape = (rows_per_rank, *first.shape[1:-1]) + if outs is not None: + for lane_size, out in zip(lane_sizes, outs, strict=True): + self._validate_output( + out, + first, + (*output_shape, lane_size), + ) + + packed_bytes = sum(value.numel() * value.element_size() for value in values) + if self._can_direct_reduce_scatter_many(): + if packed_bytes > self.max_size_bytes: + raise ValueError( + "reduce_scatter_many packed input requires " + f"{packed_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + direct_outputs = tuple( + ( + outs[index] + if outs is not None + else torch.empty( + (*output_shape, lane_size), + dtype=first.dtype, + device=first.device, + ) + ) + for index, lane_size in enumerate(lane_sizes) + ) + with self._lock: + self._check_open() + if validate_signature: + self._validate_matching_signature( + f"reduce_scatter_many:{lane_sizes}", + first, + ) + if self._direct_reduce_scatter_many(values, direct_outputs): + return direct_outputs + + if packed_bytes > _PACKED_REDUCE_SCATTER_MAX_BYTES: + # A single packed AllGather moves the same bytes as two separate + # calls but loses RCCL's smaller-message algorithm. Use the + # established per-lane path above the measured crossover; this + # keeps the convenience API from regressing large FFN gradients. + return tuple( + self.reduce_scatter( + value, + out=None if outs is None else outs[index], + validate_signature=validate_signature, + ) + for index, value in enumerate(values) + ) + + if packed_bytes > self.max_size_bytes: + raise ValueError( + "reduce_scatter_many packed input requires " + f"{packed_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + packed = torch.cat(values, dim=-1) + packed_out = torch.empty( + (packed.size(0) // self.world_size, *packed.shape[1:]), + dtype=packed.dtype, + device=packed.device, ) + with self._lock: + self._check_open() + # Include lane boundaries in the signature. Equal packed shapes + # alone do not guarantee that every rank will split the result the + # same way, which could silently associate gradients with the + # wrong lane. + if validate_signature: + self._validate_matching_signature( + f"reduce_scatter_many:{lane_sizes}", + packed, + ) + if self._direct_reduce_scatter(packed, packed_out): + pieces = tuple(packed_out.split(tuple(value.size(-1) for value in values), dim=-1)) + if outs is None: + return pieces + result: list[torch.Tensor] = [] + for piece, out in zip(pieces, outs, strict=True): + out.copy_(piece) + result.append(out) + return tuple(result) + rank_inputs = self._all_gather_transport(packed) + begin = self.rank * rows_per_rank + reduced = rank_inputs[:, begin : begin + rows_per_rank] + if not self._fused_reduction( + reduced, + packed_out, + operation="reduce_scatter", + ): + reduced = self._balanced_tree_sum(reduced) + packed_out.copy_(reduced) + + pieces = tuple(packed_out.split(tuple(value.size(-1) for value in values), dim=-1)) + if outs is None: + return pieces + result: list[torch.Tensor] = [] + for piece, out in zip(pieces, outs, strict=True): + out.copy_(piece) + result.append(out) + return tuple(result) def close(self) -> None: """Close the instance. @@ -644,6 +838,7 @@ def close(self) -> None: self._workspace = None self._validated_signatures.clear() self._closed = True + self._handle = 0 def __enter__(self) -> TorchDistributedDeterministicCollective: self._check_open() @@ -780,6 +975,25 @@ def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.T # an allocation or copy. Restrict the view to the current operation. return workspace[:required_bytes].view(input.dtype) + def _direct_all_reduce(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + + def _direct_reduce_scatter(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + + def _can_direct_reduce_scatter_many(self) -> bool: + return False + + def _direct_reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], + ) -> bool: + return False + + def _direct_all_gather(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + @staticmethod def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: world_size = rank_inputs.size(0) @@ -798,11 +1012,43 @@ def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: stride *= 2 return rank_inputs[0] + @staticmethod + def _fused_reduction( + rank_inputs: torch.Tensor, + output: torch.Tensor, + *, + operation: str, + ) -> bool: + """Use the optional ROCm fused fixed-tree kernel when available. + + The extension is deliberately optional: CPU/Gloo reference collectives + and installations built without the ROCm kernel retain the executable + Python implementation above. + """ + + if getattr(torch.version, "hip", None) is None or not rank_inputs.is_cuda: + return False + try: + from rl_engine import _C + except ImportError: + return False + if operation == "all_reduce": + fn = getattr(_C, "deterministic_collective_rocm_all_reduce", None) + if fn is not None: + fn(rank_inputs, output) + return True + elif operation == "reduce_scatter": + fn = getattr(_C, "deterministic_collective_rocm_reduce_scatter", None) + if fn is not None: + fn(rank_inputs, output) + return True + return False + class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): - """ROCm collective using RCCL AllGather strictly as tensor transport.""" + """Single-node ROCm fixed-tree collective using HIP IPC and RCCL.""" - backend_id = "rccl_all_gather_balanced_tree" + backend_id = "rocm_ipc_fixed_tree" def __init__( self, @@ -830,6 +1076,154 @@ def __init__( raise RuntimeError( "RCCL deterministic collectives require PyTorch's NCCL process-group API" ) + self._ipc_handle = 0 + self._ipc_staging: torch.Tensor | None = None + self._initialize_ipc_transport() + + @property + def workspace_size_bytes(self) -> int: + staging = self._ipc_staging + staging_bytes = 0 if staging is None else int(staging.numel()) + return staging_bytes + super().workspace_size_bytes + + def _initialize_ipc_transport(self) -> None: + if self.world_size == 1: + return + try: + from rl_engine import _C + except ImportError: + return + required_symbols = ( + "deterministic_collective_rocm_ipc_allocate", + "deterministic_collective_rocm_ipc_meta", + "deterministic_collective_rocm_ipc_create", + "deterministic_collective_rocm_ipc_synchronize", + "deterministic_collective_rocm_ipc_destroy", + "deterministic_collective_rocm_ipc_stage", + "deterministic_collective_rocm_ipc_all_reduce", + "deterministic_collective_rocm_ipc_all_reduce_input", + "deterministic_collective_rocm_ipc_reduce_scatter", + "deterministic_collective_rocm_ipc_reduce_scatter_input", + "deterministic_collective_rocm_ipc_reduce_scatter_many", + "deterministic_collective_rocm_ipc_all_gather", + "deterministic_collective_rocm_ipc_all_gather_input", + ) + if any(not hasattr(_C, symbol) for symbol in required_symbols): + return + + staging = _C.deterministic_collective_rocm_ipc_allocate(self.max_size_bytes) + handle, offset = _C.deterministic_collective_rocm_ipc_meta(staging) + local_metadata = (socket.gethostname(), handle, int(offset)) + gathered_metadata: list[tuple[str, list[int], int] | None] = [None] * self.world_size + dist.all_gather_object(gathered_metadata, local_metadata, group=self.group) + if any(metadata is None for metadata in gathered_metadata): + raise RuntimeError("failed to exchange ROCm IPC metadata") + complete_metadata = [metadata for metadata in gathered_metadata if metadata is not None] + if len({metadata[0] for metadata in complete_metadata}) != 1: + return + self._ipc_handle = int( + _C.deterministic_collective_rocm_ipc_create( + staging, + [metadata[1] for metadata in complete_metadata], + [metadata[2] for metadata in complete_metadata], + self.rank, + ) + ) + self._ipc_staging = staging + + def _direct_all_reduce(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + input_bytes = input.numel() * input.element_size() + if ( + _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES + < input_bytes + < _ROCM_IPC_SHARDED_ALL_REDUCE_MIN_BYTES + and input.numel() % self.world_size == 0 + ): + return False + + if ( + input_bytes <= _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES + or input.numel() % self.world_size != 0 + ): + _C.deterministic_collective_rocm_ipc_all_reduce_input( + handle, + input, + output, + ) + return True + + shard = self._workspace_for(input, input.numel() // self.world_size) + _C.deterministic_collective_rocm_ipc_reduce_scatter_input( + handle, + input, + shard, + ) + dist.all_gather_into_tensor(output.view(-1), shard, group=self.group) + return True + + def _direct_reduce_scatter(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_reduce_scatter_input( + handle, + input, + output, + ) + return True + + def _can_direct_reduce_scatter_many(self) -> bool: + return bool(self._ipc_handle) + + def _direct_reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], + ) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_reduce_scatter_many( + handle, + inputs, + outputs, + ) + return True + + def _direct_all_gather(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + input_bytes = input.numel() * input.element_size() + if not handle or input_bytes > _ROCM_IPC_ALL_GATHER_MAX_BYTES: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_all_gather_input( + handle, + input, + output, + ) + return True + + def close(self) -> None: + handle = getattr(self, "_ipc_handle", 0) + if handle: + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_synchronize(handle) + torch.cuda.synchronize(self.device) + self._ipc_handle = 0 + _C.deterministic_collective_rocm_ipc_destroy(handle) + self._ipc_staging = None + super().close() def create_deterministic_collective( @@ -841,9 +1235,9 @@ def create_deterministic_collective( """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. Shared caches may - replace an entry without closing it immediately because active autograd + HIP IPC or RCCL for rank-ordered transport while preserving the fixed local + reduction tree. The returned object has independent ownership. Shared caches + may replace an entry without closing it immediately because active autograd contexts can retain the previous instance until their work completes. """ diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index 7c0f2079..e793e9ce 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -641,10 +641,14 @@ def scatter(self, full: torch.Tensor) -> torch.Tensor: 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]) + # ``full`` is contiguous and split along its leading dimension, so + # each root chunk is already contiguous. Non-root ranks do not need a + # scatter list at all; avoiding those copies matters for strict CP + # output tensors, which can be large. + local_shape = (full.size(0) // self.world_size, *full.shape[1:]) + local = torch.empty(local_shape, dtype=full.dtype, device=full.device) if self.world_size == 1: - local.copy_(chunks[0]) + local.copy_(full) return local # ``src`` is a global rank even when a subgroup is supplied. @@ -660,12 +664,8 @@ def scatter(self, full: torch.Tensor) -> torch.Tensor: "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, - ) + scatter_list = list(full.chunk(self.world_size, dim=0)) if self.rank == self.root else None + dist.scatter(local, scatter_list=scatter_list, src=global_root, group=self.group) return local def reduce_scatter(self, full: torch.Tensor) -> torch.Tensor: diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index 02f97b30..453878f4 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -296,8 +296,14 @@ def forward( tp_world = tp_dist.get_world_size(group=tp_group) if tp_dist is not None else 1 gemm_tokens = rmsnorm_output_2d.size(0) * (tp_world if sequence_parallel else 1) element_size = rmsnorm_output_2d.element_size() + token_hidden_bytes = gemm_tokens * rmsnorm_output_2d.size(1) * element_size + # Sequence-parallel backward reduces the gate and up input-gradient + # lanes together. ``reduce_scatter_many`` packs those lanes along the + # final dimension, so reserve capacity for both lanes in one transport + # call rather than growing the collective (or failing) mid-backward. + reduction_bytes = token_hidden_bytes * (2 if sequence_parallel else 1) min_size_bytes = max( - gemm_tokens * rmsnorm_output_2d.size(1) * element_size, + reduction_bytes, gemm_tokens * gate_weight.size(0) * element_size, gate_weight.numel() * element_size, up_weight.numel() * element_size, @@ -467,28 +473,25 @@ def backward(ctx, grad_output: Tensor): gate_weight, disable_split_k=disable_split_k, ) - if ctx.sequence_parallel: - grad_rmsnorm_from_gate = _reduce_scatter_tokens( - grad_rmsnorm_from_gate, - tp_collective, - ) - elif tp_collective is not None: - grad_rmsnorm_from_gate = _all_reduce_inplace( - grad_rmsnorm_from_gate, - tp_collective, - ) - grad_rmsnorm_from_up = _linear_da( grad_up, up_weight, disable_split_k=disable_split_k, ) if ctx.sequence_parallel: - grad_rmsnorm_from_up = _reduce_scatter_tokens( - grad_rmsnorm_from_up, - tp_collective, + # These are independent reduction lanes. Pack them into one + # ReduceScatter while keeping each lane's balanced rank tree + # separate; adding them before the collective would change the + # floating-point parenthesization and break cross-TP bitwise + # invariance. + grad_rmsnorm_from_gate, grad_rmsnorm_from_up = tp_collective.reduce_scatter_many( + (grad_rmsnorm_from_gate, grad_rmsnorm_from_up) ) elif tp_collective is not None: + grad_rmsnorm_from_gate = _all_reduce_inplace( + grad_rmsnorm_from_gate, + tp_collective, + ) grad_rmsnorm_from_up = _all_reduce_inplace( grad_rmsnorm_from_up, tp_collective, diff --git a/setup.py b/setup.py index 3f0b3d5a..0cb6b8d0 100644 --- a/setup.py +++ b/setup.py @@ -147,6 +147,10 @@ def get_extensions(): # 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") + else: + # RCCL remains transport-only on ROCm; this HIP kernel performs + # the fixed balanced-tree arithmetic after AllGather. + cuda_sources.append("csrc/rocm/distributed/deterministic_collective.hip") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): diff --git a/tests/distributed/test_deterministic_all_gather.py b/tests/distributed/test_deterministic_all_gather.py index d6058dd3..7d1f72d0 100644 --- a/tests/distributed/test_deterministic_all_gather.py +++ b/tests/distributed/test_deterministic_all_gather.py @@ -12,7 +12,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) @@ -70,7 +70,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -95,6 +95,11 @@ def _worker(rank: int, port: int) -> None: returned = collective.all_gather(input, out=provided) assert returned is provided assert torch.equal(provided, expected) + + empty_input = torch.empty((0, 7), dtype=dtype, device=device) + empty_output = collective.all_gather(empty_input) + assert empty_output.shape == (0, 7) + assert empty_output.numel() == 0 dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py index eb406881..028429a3 100644 --- a/tests/distributed/test_deterministic_all_reduce.py +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -12,7 +12,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) @@ -61,7 +61,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -96,6 +96,25 @@ def _worker(rank: int, port: int) -> None: returned = collective.all_reduce(inplace, out=inplace) assert returned is inplace assert torch.equal(inplace, expected) + + if dtype in (torch.float16, torch.bfloat16): + packed_input = torch.cat((input, input[:1])) + packed_expected = torch.cat((expected, expected[:1])) + packed_output = collective.all_reduce(packed_input) + assert torch.equal(packed_output, packed_expected) + + output_storage = torch.empty( + packed_expected.numel() + 1, + dtype=dtype, + device=device, + ) + misaligned_output = output_storage[1:].view_as(packed_expected) + returned = collective.all_reduce( + packed_input, + out=misaligned_output, + ) + assert returned is misaligned_output + assert torch.equal(misaligned_output, packed_expected) dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_deterministic_reduce_scatter.py b/tests/distributed/test_deterministic_reduce_scatter.py index 3125f6c3..a3ee292d 100644 --- a/tests/distributed/test_deterministic_reduce_scatter.py +++ b/tests/distributed/test_deterministic_reduce_scatter.py @@ -12,7 +12,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) @@ -61,7 +61,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -98,6 +98,49 @@ def _worker(rank: int, port: int) -> None: returned = collective.reduce_scatter(input, out=provided) assert returned is provided assert torch.equal(provided, expected) + + if dtype in (torch.float16, torch.bfloat16): + output_storage = torch.empty( + expected.numel() + 1, + dtype=dtype, + device=device, + ) + misaligned_output = output_storage[1:].view_as(expected) + returned = collective.reduce_scatter( + input, + out=misaligned_output, + ) + assert returned is misaligned_output + assert torch.equal(misaligned_output, expected) + + other_generator = torch.Generator().manual_seed(20260817) + other_leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + _MAX_WORLD_SIZE * 17, + 19, + dtype=torch.float32, + generator=other_generator, + ).to(device=device, dtype=dtype) + other_leaves = list(other_leaves_tensor.unbind()) + other_input = _fixed_tree_reference( + other_leaves[start : start + leaves_per_rank] + ) + other_reduced = _fixed_tree_reference(other_leaves) + other_expected = other_reduced.chunk(tp_size, dim=0)[group_rank] + many_outs = (torch.empty_like(expected), torch.empty_like(other_expected)) + many_returned = collective.reduce_scatter_many( + (input, other_input), + outs=many_outs, + ) + assert many_returned[0] is many_outs[0] + assert many_returned[1] is many_outs[1] + assert torch.equal(many_returned[0], expected) + assert torch.equal(many_returned[1], other_expected) + many_baseline = tuple(value.clone() for value in many_returned) + for _ in range(3): + many_repeated = collective.reduce_scatter_many((input, other_input)) + assert torch.equal(many_repeated[0], many_baseline[0]) + assert torch.equal(many_repeated[1], many_baseline[1]) dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py index c550b824..5c9ffb84 100644 --- a/tests/distributed/test_transport_deterministic_collective.py +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -257,6 +257,7 @@ def test_latest_collective_api_can_skip_signature_handshakes( (peers[0], peers[0]), validate_signature=False, ) + fake_dist.peer_inputs = [torch.cat((peer, peer), dim=-1) for peer in peers] scattered = collective.reduce_scatter_many( (peers[0], peers[0]), validate_signature=False, @@ -279,6 +280,79 @@ def test_reduce_scatter_reduces_then_selects_local_leading_shard( assert torch.equal(output, expected_full.chunk(4, dim=0)[2]) +def test_reduce_scatter_many_packs_lanes_and_transports_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Each lane must retain its own balanced tree. A left fold would produce + # one for lane 0 and two for lane 1, while the fixed tree produces zero for + # both lanes in FP32. + lane_values = [ + ( + 1.0e20, + 1.0, + -1.0e20, + 1.0, + ), + ( + 1.0e20, + 2.0, + -1.0e20, + 2.0, + ), + ] + # Give each rank distinct values while preserving the cancellation pattern + # in every row. The fake transport returns these packed rank inputs. + lane_peers = [ + ( + torch.full((8, 1), lane_values[0][rank], dtype=torch.float32), + torch.full((8, 1), lane_values[1][rank], dtype=torch.float32), + ) + for rank in range(4) + ] + packed_peers = [torch.cat(lanes, dim=-1) for lanes in lane_peers] + collective, fake_dist = _make_collective(monkeypatch, packed_peers, rank=2) + + local_lanes = lane_peers[2] + outputs = (torch.empty(2, 1), torch.empty(2, 1)) + returned = collective.reduce_scatter_many(local_lanes, outs=outputs) + + assert returned[0] is outputs[0] + assert returned[1] is outputs[1] + assert torch.equal(outputs[0], torch.zeros_like(outputs[0])) + assert torch.equal(outputs[1], torch.zeros_like(outputs[1])) + assert fake_dist.tensor_transport_calls == 1 + + +def test_reduce_scatter_many_rejects_oversized_packed_input( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(4, 2, dtype=torch.float32) for _ in range(2)] + collective, fake_dist = _make_collective( + monkeypatch, + peers, + max_size_bytes=32, + ) + monkeypatch.setattr(collectives, "_PACKED_REDUCE_SCATTER_MAX_BYTES", 1024) + + with pytest.raises(ValueError, match="packed input requires"): + collective.reduce_scatter_many((peers[0], peers[0])) + assert fake_dist.tensor_transport_calls == 0 + + +def test_reduce_scatter_many_uses_separate_calls_for_large_payloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(4, 2, dtype=torch.float32) for _ in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers) + monkeypatch.setattr(collectives, "_PACKED_REDUCE_SCATTER_MAX_BYTES", 1) + + outputs = collective.reduce_scatter_many((peers[0], peers[0])) + + assert len(outputs) == 2 + assert all(torch.equal(output, torch.full_like(output, 2)) for output in outputs) + assert fake_dist.tensor_transport_calls == 2 + + def test_matching_signature_is_checked_before_tensor_transport( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/test_build_platform_collectives.py b/tests/test_build_platform_collectives.py index f094a5af..19d3a890 100644 --- a/tests/test_build_platform_collectives.py +++ b/tests/test_build_platform_collectives.py @@ -40,6 +40,7 @@ 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 "csrc/rocm/distributed/deterministic_collective.hip" 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"]