diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 58926f6429dd..1d1e72e17856 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -9,6 +9,7 @@ typedef __hip_bfloat16 nv_bfloat16; #endif +#include #include #include #include @@ -297,10 +298,12 @@ DINLINE P packed_reduce(const P* ptrs[], int idx) { template __global__ void __launch_bounds__(512, 1) - cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg, + cross_device_reduce_1stage(RankData* _dp, const RankSignals* _sg, T* __restrict__ result, int rank, int size) { using P = typename packed_t::P; using A = typename packed_t::A; + auto sg = *_sg; + auto self_sg = sg.signals[rank]; // note: we don't reorder the address so the accumulation order is the same // for all ranks, ensuring bitwise identical results auto dp = *_dp; @@ -320,12 +323,14 @@ DINLINE P* get_tmp_buf(Signal* sg) { template __global__ void __launch_bounds__(512, 1) - cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg, + cross_device_reduce_2stage(RankData* _dp, const RankSignals* _sg, T* __restrict__ result, int rank, int size) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = gridDim.x * blockDim.x; using P = typename packed_t::P; using A = typename packed_t::A; + auto sg = *_sg; + auto self_sg = sg.signals[rank]; int part = size / ngpus; int start = rank * part; int end = rank == ngpus - 1 ? size : start + part; @@ -371,12 +376,24 @@ static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t)); class CustomAllreduce { public: + enum class LifecycleState { Active, Detached, Failed }; + + struct PointerLayout { + char* base; + size_t size; + int64_t offset; + }; + + struct GraphBufferRecord { + void* local_ptr; + RankData* rank_data; + }; + int rank_; int world_size_; // Full NVLink or xGMI connection between GPUs. bool fully_connected_; - RankSignals sg_; // Stores a map from a pointer to its peer pointers from all ranks. std::unordered_map buffers_; Signal* self_sg_; @@ -397,10 +414,16 @@ class CustomAllreduce { // 3. (In Python) all gather the IPC handles. // 4. Obtain the peer pointers by opening the IPC handles, and store them in // the rank data array at corresponding positions. + RankSignals* d_rank_signals_; RankData *d_rank_data_base_, *d_rank_data_end_; std::vector graph_unreg_buffers_; + std::vector graph_buffers_; // a map from IPC handles to opened IPC pointers std::map ipc_handles_; + // Active --prepare_for_suspend--> Detached --reinit_after_resume--> + // Active. Transition work first enters Failed so errors leave the + // communicator in a terminal failed state. + LifecycleState lifecycle_state_ = LifecycleState::Active; /** * Signals are an array of ipc-enabled buffers from all ranks. @@ -419,44 +442,140 @@ class CustomAllreduce { world_size_(world_size), fully_connected_(fully_connected), self_sg_(signals[rank]), - d_rank_data_base_(reinterpret_cast(rank_data)), - d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) { + d_rank_signals_(reinterpret_cast(rank_data)), + d_rank_data_base_(reinterpret_cast(rank_data) + 1), + d_rank_data_end_(reinterpret_cast(rank_data) + + rank_data_sz / sizeof(RankData)) { + static_assert(sizeof(RankSignals) == sizeof(RankData)); + if (rank_data_sz < sizeof(RankSignals) + sizeof(RankData)) { + throw std::invalid_argument("rank data buffer is too small"); + } + RankSignals rank_signals{}; for (int i = 0; i < world_size_; i++) { - sg_.signals[i] = signals[i]; + rank_signals.signals[i] = signals[i]; + } + check_cuda(cudaMemcpy(d_rank_signals_, &rank_signals, sizeof(rank_signals), + cudaMemcpyHostToDevice), + "cudaMemcpy"); + } + + static void check_cuda(cudaError_t result, const char* operation) { + if (result != cudaSuccess) { + throw std::runtime_error(std::string(operation) + + " failed: " + cudaGetErrorString(result)); + } + } + +#if !defined(USE_ROCM) + static void check_driver(CUresult result, const char* operation) { + if (result != CUDA_SUCCESS) { + const char* error = nullptr; + cuGetErrorString(result, &error); + throw std::runtime_error( + std::string(operation) + + " failed: " + (error == nullptr ? "unknown CUDA error" : error)); + } + } +#endif + + void require_active() const { + if (lifecycle_state_ == LifecycleState::Detached) { + throw std::runtime_error("custom allreduce peer mappings are detached"); + } + if (lifecycle_state_ != LifecycleState::Active) { + throw std::runtime_error( + "custom allreduce lifecycle is in a failed state"); + } + } + + static IPC_KEY ipc_key(const void* handle) { + IPC_KEY key; + std::memcpy(key.data(), handle, key.size()); + return key; + } + + PointerLayout get_pointer_layout(void* ptr) const { +#if defined(USE_ROCM) + void* base_ptr; + if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr, + reinterpret_cast(ptr)) != + CUDA_SUCCESS) { + throw std::runtime_error("failed to get custom allreduce pointer base"); } + size_t size = 0; +#else + CUdeviceptr base; + size_t size; + check_driver( + cuMemGetAddressRange(&base, &size, reinterpret_cast(ptr)), + "cuMemGetAddressRange"); + auto base_ptr = reinterpret_cast(base); +#endif + return {reinterpret_cast(base_ptr), size, + reinterpret_cast(ptr) - reinterpret_cast(base_ptr)}; } char* open_ipc_handle(const void* ipc_handle) { - auto [it, new_handle] = - ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr}); + auto [it, new_handle] = ipc_handles_.insert({ipc_key(ipc_handle), nullptr}); if (new_handle) { char* ipc_ptr; - CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptr, - *((const cudaIpcMemHandle_t*)ipc_handle), - cudaIpcMemLazyEnablePeerAccess)); + try { + check_cuda(cudaIpcOpenMemHandle( + reinterpret_cast(&ipc_ptr), + *reinterpret_cast(ipc_handle), + cudaIpcMemLazyEnablePeerAccess), + "cudaIpcOpenMemHandle"); + } catch (...) { + ipc_handles_.erase(it); + throw; + } it->second = ipc_ptr; } return it->second; } std::pair> get_graph_buffer_ipc_meta() { + require_active(); auto num_buffers = graph_unreg_buffers_.size(); auto handle_sz = sizeof(cudaIpcMemHandle_t); std::string handles(handle_sz * num_buffers, static_cast(0)); std::vector offsets(num_buffers); for (int i = 0; i < num_buffers; i++) { auto ptr = graph_unreg_buffers_[i]; - void* base_ptr; - // note: must share the base address of each allocation, or we get wrong - // address - if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr, - (CUdeviceptr)ptr) != CUDA_SUCCESS) - throw std::runtime_error("failed to get pointer attr"); - CUDACHECK(cudaIpcGetMemHandle( - (cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr)); - offsets[i] = ((char*)ptr) - ((char*)base_ptr); + auto layout = get_pointer_layout(ptr); + check_cuda(cudaIpcGetMemHandle(reinterpret_cast( + &handles[i * handle_sz]), + layout.base), + "cudaIpcGetMemHandle"); + offsets[i] = layout.offset; + } + return std::make_pair(handles, offsets); + } + + std::pair> + get_graph_buffer_ipc_meta_for_reinit() { +#if defined(USE_ROCM) + throw std::runtime_error( + "custom allreduce suspend/resume is only supported on CUDA"); +#else + if (lifecycle_state_ != LifecycleState::Detached) { + throw std::runtime_error( + "custom allreduce must be detached before graph buffer reinit"); + } + auto handle_sz = sizeof(cudaIpcMemHandle_t); + std::string handles(handle_sz * graph_buffers_.size(), + static_cast(0)); + std::vector offsets(graph_buffers_.size()); + for (size_t i = 0; i < graph_buffers_.size(); i++) { + auto layout = get_pointer_layout(graph_buffers_[i].local_ptr); + check_cuda(cudaIpcGetMemHandle(reinterpret_cast( + &handles[i * handle_sz]), + layout.base), + "cudaIpcGetMemHandle"); + offsets[i] = layout.offset; } return std::make_pair(handles, offsets); +#endif } void check_rank_data_capacity(size_t num = 1) { @@ -466,18 +585,23 @@ class CustomAllreduce { std::to_string(d_rank_data_base_ + num - d_rank_data_end_)); } - /** - * Register already-shared IPC pointers. - */ + // Custom allreduce owns exactly one eager RankData slot. Captured inputs use + // separate graph slots below. void register_buffer(void** ptrs) { + require_active(); + if (!buffers_.empty()) { + throw std::runtime_error( + "custom allreduce supports exactly one eager registered buffer"); + } check_rank_data_capacity(); RankData data; for (int i = 0; i < world_size_; i++) { data.ptrs[i] = ptrs[i]; } auto d_data = d_rank_data_base_++; - CUDACHECK( - cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice)); + check_cuda( + cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice), + "cudaMemcpy"); buffers_[ptrs[rank_]] = d_data; } @@ -491,7 +615,19 @@ class CustomAllreduce { void register_graph_buffers( const std::vector& handles, const std::vector>& offsets) { + require_active(); auto num_buffers = graph_unreg_buffers_.size(); + if (handles.size() != world_size_ || offsets.size() != world_size_) { + throw std::runtime_error( + "custom allreduce graph metadata world size mismatch"); + } + for (int i = 0; i < world_size_; i++) { + if (handles[i].size() != num_buffers * sizeof(cudaIpcMemHandle_t) || + offsets[i].size() != num_buffers) { + throw std::runtime_error( + "custom allreduce graph metadata buffer count mismatch"); + } + } check_rank_data_capacity(num_buffers); std::vector rank_data(num_buffers); for (int i = 0; i < num_buffers; i++) { @@ -508,13 +644,126 @@ class CustomAllreduce { } } } - CUDACHECK(cudaMemcpy(d_rank_data_base_, rank_data.data(), - sizeof(RankData) * num_buffers, - cudaMemcpyHostToDevice)); + check_cuda( + cudaMemcpy(d_rank_data_base_, rank_data.data(), + sizeof(RankData) * num_buffers, cudaMemcpyHostToDevice), + "cudaMemcpy"); + for (int i = 0; i < num_buffers; i++) { + graph_buffers_.push_back( + {graph_unreg_buffers_[i], d_rank_data_base_ + i}); + } d_rank_data_base_ += num_buffers; graph_unreg_buffers_.clear(); } + // Closes all mappings in ipc_handles_, keeping the entries whose close + // failed and rethrowing the first error at the end. + void close_ipc_handles() { + std::exception_ptr close_error; + for (auto it = ipc_handles_.begin(); it != ipc_handles_.end();) { + auto result = cudaIpcCloseMemHandle(it->second); + if (result == cudaSuccess) { + it = ipc_handles_.erase(it); + } else { + if (close_error == nullptr) { + close_error = std::make_exception_ptr( + std::runtime_error(std::string("cudaIpcCloseMemHandle failed: ") + + cudaGetErrorString(result))); + } + ++it; + } + } + if (close_error != nullptr) std::rethrow_exception(close_error); + } + + void prepare_for_suspend() { +#if defined(USE_ROCM) + throw std::runtime_error( + "custom allreduce suspend/resume is only supported on CUDA"); +#else + if (lifecycle_state_ == LifecycleState::Detached) return; + require_active(); + lifecycle_state_ = LifecycleState::Failed; + check_cuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + close_ipc_handles(); + lifecycle_state_ = LifecycleState::Detached; +#endif + } + + // TODO: robust post-resume validation is deferred; the cheap invariants + // below cannot detect all stale-pointer scenarios and layout comparison is + // not sufficient either. + void reinit_after_resume(const std::vector& handles, + const std::vector>& offsets, + Signal** signals, void** buffers) { +#if defined(USE_ROCM) + throw std::runtime_error( + "custom allreduce suspend/resume is only supported on CUDA"); +#else + if (lifecycle_state_ != LifecycleState::Detached) { + throw std::runtime_error( + "custom allreduce must be detached before reinit"); + } + lifecycle_state_ = LifecycleState::Failed; + check_cuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + if (signals[rank_] != self_sg_) { + throw std::runtime_error( + "custom allreduce local signal pointer changed after resume"); + } + auto eager_buffer = buffers_.find(buffers[rank_]); + if (eager_buffer == buffers_.end()) { + throw std::runtime_error( + "custom allreduce local eager buffer changed after resume"); + } + const auto num_buffers = graph_buffers_.size(); + if (handles.size() != world_size_ || offsets.size() != world_size_) { + throw std::runtime_error( + "custom allreduce graph metadata world size mismatch"); + } + for (int i = 0; i < world_size_; i++) { + if (handles[i].size() != num_buffers * sizeof(cudaIpcMemHandle_t) || + offsets[i].size() != num_buffers) { + throw std::runtime_error( + "custom allreduce graph metadata buffer count mismatch"); + } + } + + std::vector rank_data(num_buffers); + for (size_t i = 0; i < num_buffers; i++) { + auto& rd = rank_data[i]; + for (int j = 0; j < world_size_; j++) { + if (j != rank_) { + char* handle = + open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]); + rd.ptrs[j] = handle + offsets[j][i]; + } else { + rd.ptrs[j] = graph_buffers_[i].local_ptr; + } + } + } + + RankSignals rank_signals{}; + RankData eager_rank_data{}; + for (int i = 0; i < world_size_; i++) { + rank_signals.signals[i] = signals[i]; + eager_rank_data.ptrs[i] = buffers[i]; + } + check_cuda(cudaMemcpy(d_rank_signals_, &rank_signals, sizeof(rank_signals), + cudaMemcpyHostToDevice), + "cudaMemcpy"); + check_cuda(cudaMemcpy(eager_buffer->second, &eager_rank_data, + sizeof(eager_rank_data), cudaMemcpyHostToDevice), + "cudaMemcpy"); + for (size_t i = 0; i < num_buffers; i++) { + check_cuda(cudaMemcpy(graph_buffers_[i].rank_data, &rank_data[i], + sizeof(RankData), cudaMemcpyHostToDevice), + "cudaMemcpy"); + } + check_cuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + lifecycle_state_ = LifecycleState::Active; +#endif + } + /** * Performs allreduce, assuming input has already been registered. * @@ -527,6 +776,7 @@ class CustomAllreduce { template void allreduce(cudaStream_t stream, T* input, T* output, int size, int threads = 512, int block_limit = defaultBlockLimit) { + require_active(); auto d = packed_t::P::size; if (size % d != 0) throw std::runtime_error( @@ -576,9 +826,9 @@ class CustomAllreduce { } } -#define KL(ngpus, name) \ - name<<>>(ptrs, sg_, self_sg_, output, \ - rank_, size); +#define KL(ngpus, name) \ + name<<>>(ptrs, d_rank_signals_, \ + output, rank_, size); #define REDUCE_CASE(ngpus) \ case ngpus: { \ if (force_1stage) { \ @@ -616,10 +866,14 @@ class CustomAllreduce { #undef KL } - ~CustomAllreduce() { - for (auto [_, ptr] : ipc_handles_) { - CUDACHECK(cudaIpcCloseMemHandle(ptr)); - } + void close() { + lifecycle_state_ = LifecycleState::Failed; + check_cuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + close_ipc_handles(); + } + + ~CustomAllreduce() noexcept { + for (auto [_, ptr] : ipc_handles_) cudaIpcCloseMemHandle(ptr); } }; diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/libtorch_stable/custom_all_reduce.cu index 0f7f759949a8..276f64841b6c 100644 --- a/csrc/libtorch_stable/custom_all_reduce.cu +++ b/csrc/libtorch_stable/custom_all_reduce.cu @@ -119,6 +119,11 @@ void dispose(fptr_t _fa) { delete reinterpret_cast(_fa); } +void custom_ar_close(fptr_t _fa) { + auto fa = reinterpret_cast(_fa); + fa->close(); +} + int64_t meta_size() { return sizeof(vllm::Signal); } void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { @@ -140,6 +145,14 @@ get_graph_buffer_ipc_meta(fptr_t _fa) { return std::make_tuple(bytes, offsets); } +std::tuple, std::vector> +get_graph_buffer_ipc_meta_for_reinit(fptr_t _fa) { + auto fa = reinterpret_cast(_fa); + auto [handle, offsets] = fa->get_graph_buffer_ipc_meta_for_reinit(); + std::vector bytes(handle.begin(), handle.end()); + return std::make_tuple(bytes, offsets); +} + // Use vector to represent byte data for python binding compatibility. void register_graph_buffers(fptr_t _fa, const std::vector>& handles, @@ -154,6 +167,33 @@ void register_graph_buffers(fptr_t _fa, fa->register_graph_buffers(bytes, offsets); } +void custom_ar_prepare_for_suspend(fptr_t _fa) { + auto fa = reinterpret_cast(_fa); + fa->prepare_for_suspend(); +} + +void custom_ar_reinit_after_resume( + fptr_t _fa, const std::vector>& handles, + const std::vector>& offsets, + const std::vector& fake_signal_ptrs, + const std::vector& fake_buffer_ptrs) { + auto fa = reinterpret_cast(_fa); + STD_TORCH_CHECK(fake_signal_ptrs.size() == fa->world_size_); + STD_TORCH_CHECK(fake_buffer_ptrs.size() == fa->world_size_); + std::vector bytes; + bytes.reserve(handles.size()); + for (const auto& handle : handles) { + bytes.emplace_back(handle.begin(), handle.end()); + } + vllm::Signal* signal_ptrs[8]; + void* buffer_ptrs[8]; + for (int i = 0; i < fa->world_size_; i++) { + signal_ptrs[i] = reinterpret_cast(fake_signal_ptrs[i]); + buffer_ptrs[i] = reinterpret_cast(fake_buffer_ptrs[i]); + } + fa->reinit_after_resume(bytes, offsets, signal_ptrs, buffer_ptrs); +} + std::tuple allocate_shared_buffer_and_handle( int64_t size) { int device_index; @@ -197,6 +237,21 @@ fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) { return reinterpret_cast(ipc_ptr); } +torch::stable::Tensor get_mem_handle(fptr_t buffer) { + auto handle = torch::stable::empty( + {static_cast(sizeof(cudaIpcMemHandle_t))}, + torch::headeronly::ScalarType::Byte, std::nullopt, + torch::stable::Device(torch::stable::DeviceType::CPU)); + STD_CUDA_CHECK(cudaIpcGetMemHandle( + reinterpret_cast(handle.mutable_data_ptr()), + reinterpret_cast(buffer))); + return handle; +} + +void close_mem_handle(fptr_t ptr) { + STD_CUDA_CHECK(cudaIpcCloseMemHandle(reinterpret_cast(ptr))); +} + void free_shared_buffer(fptr_t buffer) { STD_CUDA_CHECK(cudaFree(reinterpret_cast(buffer))); } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 7cf34d8b03ab..3ff880920437 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -374,16 +374,27 @@ void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, torch::stable::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); void dispose(fptr_t _fa); +void custom_ar_close(fptr_t _fa); int64_t meta_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); std::tuple, std::vector> get_graph_buffer_ipc_meta(fptr_t _fa); +std::tuple, std::vector> +get_graph_buffer_ipc_meta_for_reinit(fptr_t _fa); void register_graph_buffers(fptr_t _fa, const std::vector>& handles, const std::vector>& offsets); +void custom_ar_prepare_for_suspend(fptr_t _fa); +void custom_ar_reinit_after_resume( + fptr_t _fa, const std::vector>& handles, + const std::vector>& offsets, + const std::vector& fake_signal_ptrs, + const std::vector& fake_buffer_ptrs); std::tuple allocate_shared_buffer_and_handle( int64_t size); int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +torch::stable::Tensor get_mem_handle(int64_t buffer); +void close_mem_handle(int64_t ptr); void free_shared_buffer(int64_t buffer); // Activation kernels (shared CUDA/ROCm) diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 158999a6633b..1bbe572b6f92 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -888,13 +888,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " "int reg_buffer_sz_bytes) -> ()"); custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("custom_ar_close(int fa) -> ()"); custom_ar.def("meta_size() -> int"); custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "get_graph_buffer_ipc_meta_for_reinit(int fa) -> (int[], int[])"); custom_ar.def( "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("custom_ar_prepare_for_suspend(int fa) -> ()"); + custom_ar.def( + "custom_ar_reinit_after_resume(int fa, int[][] handles, " + "int[][] offsets, int[] signal_ptrs, int[] buffer_ptrs) -> ()"); custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("get_mem_handle(int ptr) -> Tensor"); + custom_ar.def("close_mem_handle(int ptr) -> ()"); custom_ar.def("free_shared_buffer(int ptr) -> ()"); } @@ -909,13 +918,22 @@ STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("custom_ar_close", TORCH_BOX(&custom_ar_close)); custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); custom_ar.impl("get_graph_buffer_ipc_meta", TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("get_graph_buffer_ipc_meta_for_reinit", + TORCH_BOX(&get_graph_buffer_ipc_meta_for_reinit)); custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("custom_ar_prepare_for_suspend", + TORCH_BOX(&custom_ar_prepare_for_suspend)); + custom_ar.impl("custom_ar_reinit_after_resume", + TORCH_BOX(&custom_ar_reinit_after_resume)); custom_ar.impl("allocate_shared_buffer_and_handle", TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("get_mem_handle", TORCH_BOX(&get_mem_handle)); + custom_ar.impl("close_mem_handle", TORCH_BOX(&close_mem_handle)); custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); } diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index edddb6ec8455..8f1e9196ffe1 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -8,7 +8,9 @@ import torch import torch.distributed as dist +from vllm import _custom_ops as ops from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from vllm.distributed.device_communicators.custom_all_reduce import _LifecycleState from vllm.distributed.parallel_state import get_tp_group, graph_capture from ..utils import ( @@ -104,22 +106,135 @@ def eager_allreduce( num_communication = rank // tp_size + 1 sz = 1024 fa = get_tp_group().device_communicator.ca_comm + assert fa is not None + with pytest.raises(RuntimeError, match="exactly one eager registered buffer"): + ops.register_buffer(fa._ptr, fa.buffer_ptrs) inp = torch.ones(sz, dtype=torch.float32, device=device) out = inp for _ in range(num_communication): out = fa.all_reduce(out, registered=False) torch.testing.assert_close(out, inp * (tp_size**num_communication)) + if torch.version.hip is None: + expected_peer_mappings = 2 * (fa.world_size - 1) + assert len(fa._open_peer_ptrs) == expected_peer_mappings + fa.prepare_for_suspend() + assert not fa._open_peer_ptrs + fa.prepare_for_suspend() + with pytest.raises(RuntimeError, match="detached"): + tensor_model_parallel_all_reduce(inp) + fa.reinit_after_resume() + assert len(fa._open_peer_ptrs) == expected_peer_mappings + fa.reinit_after_resume() + eager_out = tensor_model_parallel_all_reduce(inp) + torch.testing.assert_close(eager_out, inp * tp_size) + inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device) out = inp for _ in range(num_communication): out = fa.all_reduce(out, registered=False) torch.testing.assert_close(out, inp * (tp_size**num_communication)) + if torch.version.hip is None: + fa.prepare_for_suspend() + fa.close() + fa.close() + + +@ray.remote(num_gpus=1, max_calls=1) +def graph_allreduce_suspend_resume( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +): + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + m.delenv("HIP_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + tp_group = get_tp_group() + device_communicator = tp_group.device_communicator + assert device_communicator is not None + fa = device_communicator.ca_comm + assert fa is not None and not fa.disabled + + selected_calls = 0 + original_custom_all_reduce = fa.custom_all_reduce + + def tracked_custom_all_reduce(input_: torch.Tensor) -> torch.Tensor | None: + nonlocal selected_calls + selected_calls += 1 + return original_custom_all_reduce(input_) + + m.setattr(fa, "custom_all_reduce", tracked_custom_all_reduce) + inp1 = torch.full((1024,), fa.rank + 1, dtype=torch.float32, device=device) + inp2 = torch.full((2048,), fa.rank + 2, dtype=torch.float32, device=device) + with graph_capture(device=device) as graph_capture_context: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + out1 = tensor_model_parallel_all_reduce(inp1) + out1_repeat = tensor_model_parallel_all_reduce(inp1) + out2 = tensor_model_parallel_all_reduce(inp2) + assert selected_calls == 3 + assert fa._lifecycle_state == _LifecycleState.ACTIVE + + expected_peer_mappings = 2 * (fa.world_size - 1) + assert len(fa._open_peer_ptrs) == expected_peer_mappings + fa.prepare_for_suspend() + assert not fa._open_peer_ptrs + fa.prepare_for_suspend() + with pytest.raises(RuntimeError, match="detached"): + tensor_model_parallel_all_reduce(inp1) + fa.reinit_after_resume() + assert len(fa._open_peer_ptrs) == expected_peer_mappings + fa.reinit_after_resume() + + for replay_idx in range(2): + inp1.fill_(fa.rank + 3 + replay_idx) + inp2.fill_(fa.rank + 4 + replay_idx) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close( + out1, + torch.full_like( + out1, sum(range(3 + replay_idx, tp_size + 3 + replay_idx)) + ), + ) + torch.testing.assert_close(out1_repeat, out1) + torch.testing.assert_close( + out2, + torch.full_like( + out2, sum(range(4 + replay_idx, tp_size + 4 + replay_idx)) + ), + ) + + selected_before_eager = selected_calls + eager_out = tensor_model_parallel_all_reduce(inp1) + assert selected_calls == selected_before_eager + 1 + torch.testing.assert_close(eager_out, out1) + fa.close() + fa.close() @pytest.mark.parametrize("tp_size", [2]) @pytest.mark.parametrize("pipeline_parallel_size", [1, 2]) -@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) +@pytest.mark.parametrize( + "test_target", + [ + eager_allreduce, + graph_allreduce, + pytest.param( + graph_allreduce_suspend_resume, + marks=pytest.mark.skipif( + torch.version.hip is not None, + reason="Custom allreduce suspend/resume requires CUDA.", + ), + ), + ], +) def test_custom_allreduce( monkeypatch: pytest.MonkeyPatch, tp_size, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index a609035560e6..56f7e2eb6c30 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2907,6 +2907,10 @@ def dispose(fa: int) -> None: torch.ops._C_custom_ar.dispose(fa) +def custom_ar_close(fa: int) -> None: + torch.ops._C_custom_ar.custom_ar_close(fa) + + def meta_size() -> int: return torch.ops._C_custom_ar.meta_size() @@ -2919,12 +2923,34 @@ def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]: return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa) +def get_graph_buffer_ipc_meta_for_reinit( + fa: int, +) -> tuple[list[int], list[int]]: + return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta_for_reinit(fa) + + def register_graph_buffers( fa: int, handles: list[list[int]], offsets: list[list[int]] ) -> None: torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets) +def custom_ar_prepare_for_suspend(fa: int) -> None: + torch.ops._C_custom_ar.custom_ar_prepare_for_suspend(fa) + + +def custom_ar_reinit_after_resume( + fa: int, + handles: list[list[int]], + offsets: list[list[int]], + signal_ptrs: list[int], + buffer_ptrs: list[int], +) -> None: + torch.ops._C_custom_ar.custom_ar_reinit_after_resume( + fa, handles, offsets, signal_ptrs, buffer_ptrs + ) + + def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]: return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size) @@ -2933,6 +2959,14 @@ def open_mem_handle(mem_handle: torch.Tensor): return torch.ops._C_custom_ar.open_mem_handle(mem_handle) +def get_mem_handle(ptr: int) -> torch.Tensor: + return torch.ops._C_custom_ar.get_mem_handle(ptr) + + +def close_mem_handle(ptr: int) -> None: + torch.ops._C_custom_ar.close_mem_handle(ptr) + + def free_shared_buffer(ptr: int) -> None: torch.ops._C_custom_ar.free_shared_buffer(ptr) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b92015b18808..db484da10268 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -508,6 +508,12 @@ def destroy(self): self.pynccl_comm.destroy() self.pynccl_comm = None if self.ca_comm is not None: + try: + self.ca_comm.close() + except Exception: + logger.warning( + "Failed to close custom allreduce communicator.", exc_info=True + ) self.ca_comm = None if self.fi_ar_comm is not None: self.fi_ar_comm.destroy() diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 95db6cc92459..74420347a2cb 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import contextmanager +import enum +from contextlib import contextmanager, suppress from typing import cast import torch @@ -28,6 +29,13 @@ logger = init_logger(__name__) +class _LifecycleState(enum.Enum): + ACTIVE = enum.auto() + DETACHED = enum.auto() + FAILED = enum.auto() + CLOSED = enum.auto() + + def _can_p2p(rank: int, world_size: int) -> bool: for i in range(world_size): if i == rank: @@ -71,6 +79,11 @@ def __init__( """ self._IS_CAPTURING = False self.disabled = True + # Attributes needed by close() when __init__ exits early or fails. + self._ptr = 0 + self._lifecycle_state = _LifecycleState.ACTIVE + self._open_peer_ptrs: set[int] = set() + self._owned_ptrs: set[int] = set() if not custom_ar: # disable because of missing custom allreduce library @@ -166,32 +179,46 @@ def __init__( ) return - self.disabled = False - # Buffers memory are owned by this Python class and passed to C++. - # Metadata composes of two parts: metadata for synchronization and a - # temporary buffer for storing intermediate allreduce results. - self.meta_ptrs = self.create_shared_buffer( - ops.meta_size() + max_size, group=group, uncached=True - ) - # This is a pre-registered IPC buffer. In eager mode, input tensors - # are first copied into this buffer before allreduce is performed - self.buffer_ptrs = self.create_shared_buffer(max_size, group=group) - # This is a buffer for storing the tuples of pointers pointing to - # IPC buffers from all ranks. Each registered tuple has size of - # 8*world_size bytes where world_size is at most 8. Allocating 8MB - # is enough for 131072 such tuples. The largest model I've seen only - # needs less than 10000 of registered tuples. - self.rank_data = torch.empty( - 8 * 1024 * 1024, dtype=torch.uint8, device=self.device - ) self.max_size = max_size self.rank = rank self.world_size = world_size self.fully_connected = fully_connected - self._ptr = ops.init_custom_ar( - self.meta_ptrs, self.rank_data, rank, self.fully_connected - ) - ops.register_buffer(self._ptr, self.buffer_ptrs) + try: + # Buffers memory are owned by this Python class and passed to C++. + # Metadata composes of two parts: metadata for synchronization + # and a temporary buffer for storing intermediate allreduce + # results. + self.meta_ptrs = self.create_shared_buffer( + ops.meta_size() + max_size, group=group, uncached=True + ) + self._owned_ptrs.add(self.meta_ptrs[rank]) + self._open_peer_ptrs.update( + ptr for i, ptr in enumerate(self.meta_ptrs) if i != rank + ) + # This is a pre-registered IPC buffer. In eager mode, input + # tensors are first copied into this buffer before allreduce is + # performed + self.buffer_ptrs = self.create_shared_buffer(max_size, group=group) + self._owned_ptrs.add(self.buffer_ptrs[rank]) + self._open_peer_ptrs.update( + ptr for i, ptr in enumerate(self.buffer_ptrs) if i != rank + ) + # This is a buffer for storing the tuples of pointers pointing to + # IPC buffers from all ranks. The first RankData-sized slot holds + # the stable RankSignals pointer table; the remaining slots hold + # eager and captured-graph RankData records. + self.rank_data = torch.empty( + 8 * 1024 * 1024, dtype=torch.uint8, device=self.device + ) + self._ptr = ops.init_custom_ar( + self.meta_ptrs, self.rank_data, rank, self.fully_connected + ) + ops.register_buffer(self._ptr, self.buffer_ptrs) + except Exception: + with suppress(Exception): + self.close() + raise + self.disabled = False @contextmanager def capture(self): @@ -200,6 +227,7 @@ def capture(self): `register_graph_buffers` call at the end of the context. It records all the buffer addresses used in the CUDA graph. """ + self._require_active() try: self._IS_CAPTURING = True yield @@ -211,6 +239,13 @@ def capture(self): def register_graph_buffers(self): handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) logger.info("Registering %d cuda graph addresses", len(offset)) + handles, offsets = self._exchange_int_lists(handle, offset) + ops.register_graph_buffers(self._ptr, handles, offsets) + + def _exchange_int_lists( + self, handle: list[int], offset: list[int] + ) -> tuple[list[list[int]], list[list[int]]]: + """Exchange a pair of int lists with every rank in the group.""" # We cannot directly use `dist.all_gather_object` here # because it is incompatible with `gloo` backend under inference mode. # see https://github.com/pytorch/pytorch/issues/126032 for details. @@ -225,7 +260,13 @@ def register_graph_buffers(self): # Unpack list of tuples to tuple of lists. handles = cast(list[list[int]], [d[0] for d in all_data]) offsets = cast(list[list[int]], [d[1] for d in all_data]) - ops.register_graph_buffers(self._ptr, handles, offsets) + return handles, offsets + + def _require_active(self) -> None: + if self._lifecycle_state == _LifecycleState.DETACHED: + raise RuntimeError("Custom allreduce peer mappings are detached.") + if self._lifecycle_state != _LifecycleState.ACTIVE: + raise RuntimeError("Custom allreduce lifecycle is not active.") def should_custom_ar(self, inp: torch.Tensor): if self.disabled: @@ -251,6 +292,7 @@ def all_reduce( IPC-registered. Otherwise, inp is first copied into a pre-registered buffer. """ + self._require_active() if out is None: out = torch.empty_like(inp) if registered: @@ -264,7 +306,10 @@ def all_reduce( def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None: """The main allreduce API that provides support for cuda graph.""" # When custom allreduce is disabled, this will be None. - if self.disabled or not self.should_custom_ar(input): + if self.disabled: + return None + self._require_active() + if not self.should_custom_ar(input): return None if self._IS_CAPTURING: if torch.cuda.is_current_stream_capturing(): @@ -279,16 +324,127 @@ def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None: # latency) compared to the performance gain of using custom kernels return self.all_reduce(input, registered=False) - def close(self): - if not self.disabled and self._ptr: - if ops is not None: + def prepare_for_suspend(self) -> None: + """Close every imported peer mapping ahead of GPU suspend. + + Failure leaves the communicator in a terminal failed state. + Repeated successful calls are no-ops. + """ + if self.disabled or self._lifecycle_state == _LifecycleState.DETACHED: + return + self._require_active() + self._lifecycle_state = _LifecycleState.FAILED + if not current_platform.is_cuda(): + raise RuntimeError( + "Custom allreduce suspend/resume is only supported on CUDA." + ) + torch.cuda.synchronize(self.device) + ops.custom_ar_prepare_for_suspend(self._ptr) + close_error: Exception | None = None + for peer_ptr in tuple(self._open_peer_ptrs): + try: + ops.close_mem_handle(peer_ptr) + except Exception as exc: + if close_error is None: + close_error = exc + else: + self._open_peer_ptrs.remove(peer_ptr) + if close_error is not None: + raise close_error + self._lifecycle_state = _LifecycleState.DETACHED + + def reinit_after_resume(self) -> None: + """Reopen peer mappings and rewrite device pointer tables after resume. + + Failure leaves the communicator in a terminal failed state. + Repeated successful calls are no-ops. + """ + if self.disabled or self._lifecycle_state == _LifecycleState.ACTIVE: + return + if self._lifecycle_state != _LifecycleState.DETACHED: + raise RuntimeError("Custom allreduce must be detached before reinit.") + self._lifecycle_state = _LifecycleState.FAILED + if not current_platform.is_cuda(): + raise RuntimeError( + "Custom allreduce suspend/resume is only supported on CUDA." + ) + torch.cuda.synchronize(self.device) + # Re-export and exchange the owner handles of the two shared buffers. + meta_handle = ops.get_mem_handle(self.meta_ptrs[self.rank]).tolist() + buffer_handle = ops.get_mem_handle(self.buffer_ptrs[self.rank]).tolist() + meta_handles, buffer_handles = self._exchange_int_lists( + meta_handle, buffer_handle + ) + for ptrs, peer_handles in ( + (self.meta_ptrs, meta_handles), + (self.buffer_ptrs, buffer_handles), + ): + for peer_rank in range(self.world_size): + if peer_rank == self.rank: + continue + peer_ptr = ops.open_mem_handle( + torch.tensor(peer_handles[peer_rank], dtype=torch.uint8) + ) + self._open_peer_ptrs.add(peer_ptr) + ptrs[peer_rank] = peer_ptr + handle, offset = ops.get_graph_buffer_ipc_meta_for_reinit(self._ptr) + handles, offsets = self._exchange_int_lists(handle, offset) + ops.custom_ar_reinit_after_resume( + self._ptr, handles, offsets, self.meta_ptrs, self.buffer_ptrs + ) + torch.cuda.synchronize(self.device) + self._lifecycle_state = _LifecycleState.ACTIVE + + def close(self) -> None: + """Release all custom-allreduce resources, surfacing failures. + + A successful close is terminal; repeated calls are no-ops. + """ + if self._lifecycle_state == _LifecycleState.CLOSED or ops is None: + return + if self._ptr == 0 and not self._open_peer_ptrs and not self._owned_ptrs: + self.disabled = True + self._lifecycle_state = _LifecycleState.CLOSED + return + self._lifecycle_state = _LifecycleState.FAILED + first_error: Exception | None = None + native_closed = self._ptr == 0 + if self._ptr: + try: + ops.custom_ar_close(self._ptr) ops.dispose(self._ptr) - self._ptr = 0 - self.free_shared_buffer(self.meta_ptrs, rank=self.rank) - self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) + self._ptr = 0 + native_closed = True + except Exception as exc: + first_error = exc + for peer_ptr in tuple(self._open_peer_ptrs): + try: + ops.close_mem_handle(peer_ptr) + except Exception as exc: + if first_error is None: + first_error = exc + else: + self._open_peer_ptrs.remove(peer_ptr) + # Fail closed: keep owned buffers alive if the native close failed. + if native_closed: + for owner_ptr in tuple(self._owned_ptrs): + try: + ops.free_shared_buffer(owner_ptr) + except Exception as exc: + if first_error is None: + first_error = exc + else: + self._owned_ptrs.remove(owner_ptr) + if first_error is not None: + raise RuntimeError( + "Failed to close custom allreduce resources." + ) from first_error + self.disabled = True + self._lifecycle_state = _LifecycleState.CLOSED def __del__(self): - self.close() + with suppress(Exception): + self.close() @staticmethod def create_shared_buffer(