Commit 71efc559a for llama.cpp
commit 71efc559a00b236851f6ba00a629d776f2416d2b
Author: zyang-dev <267119621+zyang-dev@users.noreply.github.com>
Date: Tue Nov 3 01:26:19 2026 -0800
Add nccl support for multi-node tensor parallelism (#28967)
* Add a backend interface for collectives across ranks running in separate processes, implement it in the CUDA backend using NCCL, and use it in the RPC server when available.
* Drop ggml_backend_comm_allreduce_rank_t and reuse ggml_backend_comm_allreduce_tensor for per-rank communicators
* reuse ggml_backend_cuda_comm_allreduce_nccl for per-rank communicators. Remove the helper function backend_comm_proc_address
* Change ggml_backend_comm_get_unique_id_t and ggml_backend_comm_init_rank_t to more generic multi-process backend communicator interfaces
* set maxCTAs in ncclConfig_t to 8 for dgx spark in ggml_backend_cuda_comm_init_rank
diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h
index 735ed2c2b..afb86338a 100644
--- a/ggml/include/ggml-backend.h
+++ b/ggml/include/ggml-backend.h
@@ -211,6 +211,18 @@ extern "C" {
// Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+
+ // Builds the communicator for the case where each rank is a separate process, so a process
+ // owns one backend of the communicator rather than all of them.
+ //
+ // The caller provides bootstrap_callback, which moves bytes between the ranks before the
+ // communicator exists. Each rank passes size bytes in send_buf, and recv_buf receives
+ // world_size * size bytes in rank order. The backend decides what to exchange and how many
+ // times to call it, and every rank must make the same calls in the same order.
+ // comm_init_rank blocks until all world_size ranks have joined.
+ typedef bool (*ggml_backend_comm_bootstrap_callback)(void * user_data, const void * send_buf, void * recv_buf, size_t size);
+ typedef void * (*ggml_backend_comm_init_rank_t)(ggml_backend_t backend, int rank, int world_size, ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data);
+
typedef void (*ggml_backend_comm_free_t)(void * comm_ctx);
typedef bool (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 06956ce7a..84fc202a1 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1117,6 +1117,18 @@ struct ggml_backend_cuda_comm_context {
ggml_cuda_ar_pipeline * ar_pipeline = nullptr;
+ // Number of ranks in the communicator when each rank is a separate process. Zero unless
+ // comm_init_rank created this context.
+ int world_size = 0;
+
+ // Number of ranks taking part in the reduction. For a communicator from comm_init_rank this is
+ // world_size. Otherwise it is the number of backends this process holds. The NCCL reduction
+ // uses it to choose the tensor size below which it reduces in F32 rather than converting to
+ // BF16 first.
+ size_t n_ranks() const {
+ return world_size > 0 ? (size_t) world_size : backends.size();
+ }
+
#ifdef GGML_USE_NCCL
std::vector<ncclComm_t> comms;
#endif // GGML_USE_NCCL
@@ -1144,6 +1156,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
}
const size_t n_backends = comm_ctx->backends.size();
+ const size_t n_ranks = comm_ctx->n_ranks();
for (size_t i = 0; i < n_backends; ++i) {
GGML_ASSERT(tensors[i] != nullptr);
@@ -1153,7 +1166,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
// For small tensors, simply reduce them as FP32.
// The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0.
- if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) {
+ if ((n_ranks <= 2 && ne < 32768) || (n_ranks == 3 && ne < 131072) || (n_ranks >= 4 && ne < 262144)) {
for (size_t i = 0; i < n_backends; ++i) {
if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context;
@@ -1384,6 +1397,90 @@ static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_ba
return ret;
}
+#ifdef GGML_USE_NCCL
+// Builds the communicator with ncclCommInitRank, for the case where each rank is its own
+// process. Blocks until all world_size ranks have joined.
+static void * ggml_backend_cuda_comm_init_rank(ggml_backend_t backend, int rank, int world_size,
+ ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data) {
+ if (world_size < 1 || rank < 0 || rank >= world_size) {
+ GGML_LOG_ERROR("%s: invalid rank %d for world size %d\n", __func__, rank, world_size);
+ return nullptr;
+ }
+
+ // Every rank sends whether its local checks passed, and rank 0 also sends the NCCL unique id.
+ // So all ranks either join together or fail together, and no rank waits in ncclCommInitRank alone.
+ struct bootstrap_msg {
+ uint8_t ok;
+ ncclUniqueId id;
+ };
+ bootstrap_msg local;
+ memset(&local, 0, sizeof(local));
+ local.ok = 1;
+
+ if (!ggml_backend_is_cuda(backend)) {
+ GGML_LOG_ERROR("%s: backend is not a CUDA backend\n", __func__);
+ local.ok = 0;
+ } else {
+ // NCCL requires one distinct physical GPU per rank, see ggml_backend_cuda_comm_init_nccl.
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ if (info.device_count > info.physical_device_count) {
+ GGML_LOG_ERROR("%s: CUDA virtual devices are in use, which NCCL does not support\n", __func__);
+ local.ok = 0;
+ }
+ }
+ if (local.ok && rank == 0) {
+ const ncclResult_t rc = ncclGetUniqueId(&local.id);
+ if (rc != ncclSuccess) {
+ GGML_LOG_ERROR("%s: ncclGetUniqueId failed (%s)\n", __func__, ncclGetErrorString(rc));
+ local.ok = 0;
+ }
+ }
+
+ std::vector<bootstrap_msg> all(world_size);
+ if (!bootstrap_callback(bootstrap_callback_data, &local, all.data(), sizeof(local))) {
+ GGML_LOG_ERROR("%s: bootstrap exchange failed on rank %d\n", __func__, rank);
+ return nullptr;
+ }
+ for (int i = 0; i < world_size; i++) {
+ if (!all[i].ok) {
+ GGML_LOG_ERROR("%s: rank %d failed its setup checks\n", __func__, i);
+ return nullptr;
+ }
+ }
+
+ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+
+ auto * ret = new ggml_backend_cuda_comm_context;
+ ret->backends.push_back(backend);
+ ret->dev_ids.push_back(cuda_ctx->device);
+ ret->world_size = world_size;
+ ret->comms.resize(1);
+ ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl;
+
+ // On DGX Spark, NCCL picks 64 channels by default, which is slower and uses more memory than 8.
+ // NCCL_MAX_CTAS still overrides this.
+ ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+ const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+ if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+ config.maxCTAs = 8;
+ }
+
+ ggml_cuda_set_device(cuda_ctx->device);
+ const ncclResult_t rc = ncclCommInitRankConfig(&ret->comms[0], world_size, all[0].id, rank, &config);
+ if (rc != ncclSuccess) {
+ GGML_LOG_ERROR("%s: ncclCommInitRankConfig failed for rank %d/%d (%s)\n",
+ __func__, rank, world_size, ncclGetErrorString(rc));
+ ret->comms.clear();
+ delete ret;
+ return nullptr;
+ }
+
+ GGML_LOG_INFO("%s: joined NCCL communicator as rank %d/%d on device %d\n",
+ __func__, rank, world_size, cuda_ctx->device);
+ return ret;
+}
+#endif // GGML_USE_NCCL
+
// Top-level dispatch -- calls the function pointer chosen by comm_init.
// Returns false to let the meta-backend's butterfly run.
static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) {
@@ -6384,6 +6481,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
return (void *)ggml_backend_cuda_comm_allreduce_tensor;
}
+#ifdef GGML_USE_NCCL
+ if (strcmp(name, "ggml_backend_comm_init_rank") == 0) {
+ return (void *)ggml_backend_cuda_comm_init_rank;
+ }
+#endif // GGML_USE_NCCL
if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
return (void *)ggml_backend_cuda_register_host_buffer;
}
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
index 5cf9339bd..966bd9975 100644
--- a/ggml/src/ggml-rpc/ggml-rpc.cpp
+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -1430,6 +1430,9 @@ public:
private:
void sync_all_backends();
+ void comm_free_device(uint32_t dev_id);
+ static bool comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors);
+ static bool comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size);
bool get_cached_file(uint64_t hash, std::vector<uint8_t> & data);
ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor);
ggml_tensor * create_node(uint64_t id,
@@ -1438,15 +1441,26 @@ private:
std::unordered_map<uint64_t, struct ggml_tensor*> & tensor_map);
- // pairwise allreduce over a direct connection to the peer server
+ // Per-device state for one communicator.
struct comm_state {
- socket_ptr peer;
+ // Sockets to the other servers. Rank 0 holds one per peer in rank order, so peers[i] is
+ // rank i + 1. Every other rank holds one to rank 0. The pairwise path reduces over peers[0].
+ std::vector<socket_ptr> peers;
uint32_t rank = 0;
uint32_t world = 0;
+ ggml_backend_t pairwise_backend = nullptr;
ggml_backend_buffer_ptr scratch;
size_t scratch_size = 0;
std::vector<uint8_t> send_buf;
std::vector<uint8_t> recv_buf;
+
+ // The backend's communicator. Null on the pairwise path, which has none.
+ void * backend_comm_ctx = nullptr;
+
+ // The chosen reduction, set once a path is agreed. Receives backend_comm_ctx, or
+ // comm_state on the pairwise path. A communicator covers one device, so the caller
+ // passes a one-element tensor array.
+ ggml_backend_comm_allreduce_tensor_t allreduce_fn = nullptr;
};
std::vector<ggml_backend_t> backends;
@@ -2225,78 +2239,188 @@ void rpc_server::sync_all_backends() {
}
}
-// The comm link between two servers uses the same caps negotiation as the client HELLO,
+void rpc_server::comm_free_device(uint32_t dev_id) {
+ comm_state & state = comm_states[dev_id];
+ if (state.backend_comm_ctx != nullptr) {
+ auto comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address(
+ ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id])), "ggml_backend_comm_free");
+ if (comm_free != nullptr) {
+ comm_free(state.backend_comm_ctx);
+ }
+ }
+ state = comm_state();
+}
+
+// Exchange RPC transport capabilities over a peer connection. recv_first selects which side
+// receives first. Both sides receiving first would deadlock.
+static bool comm_caps_handshake(const socket_ptr & peer, bool recv_first) {
+ uint8_t local_caps[RPC_CONN_CAPS_SIZE] = {};
+ uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+ peer->get_caps(local_caps);
+ if (recv_first) {
+ if (!peer->recv_data(remote_caps, sizeof(remote_caps)) ||
+ !peer->send_data(local_caps, sizeof(local_caps))) {
+ return false;
+ }
+ } else {
+ if (!peer->send_data(local_caps, sizeof(local_caps)) ||
+ !peer->recv_data(remote_caps, sizeof(remote_caps))) {
+ return false;
+ }
+ }
+ peer->update_caps(remote_caps);
+ return true;
+}
+
+// The comm link between servers uses the same caps negotiation as the client HELLO,
// so it gets the same transport upgrades (e.g. RDMA).
bool rpc_server::comm_init(const rpc_msg_comm_init_req & request, rpc_msg_comm_init_rsp & response) {
response.ok = 0;
- if (request.device >= backends.size() || request.world != 2 || request.rank >= request.world) {
+ if (request.device >= backends.size() || request.world < 2 || request.rank >= request.world) {
return true;
}
- comm_state & state = comm_states[request.device];
- if (state.peer != nullptr) {
+ const uint32_t dev_id = request.device;
+ comm_state & state = comm_states[dev_id];
+ if (state.allreduce_fn != nullptr) {
response.ok = 1;
return true;
}
- uint8_t local_caps[RPC_CONN_CAPS_SIZE] = {};
- uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+
+ // Joining a communicator blocks until every rank arrives, so drain outstanding work first.
+ sync_all_backends();
+
+ // Rank 0 listens and accepts world-1 peers; every other rank connects back to it.
if (request.rank == 0) {
socket_ptr srv = socket_t::create_server("0.0.0.0", request.port);
if (srv == nullptr) {
LOG_ERROR("[%s] failed to listen on comm port %u\n", __func__, request.port);
return true;
}
- state.peer = srv->accept();
- if (state.peer == nullptr) {
- return true;
- }
- if (!state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
- state.peer = nullptr;
- return true;
- }
- state.peer->get_caps(local_caps);
- if (!state.peer->send_data(local_caps, sizeof(local_caps))) {
- state.peer = nullptr;
- return true;
+ // Store peers[i] as rank i + 1, as comm_bootstrap_exchange needs.
+ // Each peer sends its rank because peers may connect in any order.
+ state.peers.resize(request.world - 1);
+ for (uint32_t i = 1; i < request.world; i++) {
+ socket_ptr peer = srv->accept();
+ uint32_t peer_rank = 0;
+ if (peer == nullptr || !comm_caps_handshake(peer, /*recv_first =*/ true) ||
+ !peer->recv_data(&peer_rank, sizeof(peer_rank)) ||
+ peer_rank == 0 || peer_rank >= request.world || state.peers[peer_rank - 1] != nullptr) {
+ LOG_ERROR("[%s] bad peer connection (rank %u)\n", __func__, peer_rank);
+ state.peers.clear();
+ return true;
+ }
+ state.peers[peer_rank - 1] = std::move(peer);
}
- state.peer->update_caps(remote_caps);
} else {
const std::string host(request.host, strnlen(request.host, sizeof(request.host)));
+ socket_ptr peer;
// rank 0 may not be listening yet, retry for a few seconds
- for (int i = 0; i < 100 && state.peer == nullptr; i++) {
- state.peer = socket_t::connect(host.c_str(), request.port);
- if (state.peer == nullptr) {
+ for (int i = 0; i < 100 && peer == nullptr; i++) {
+ peer = socket_t::connect(host.c_str(), request.port);
+ if (peer == nullptr) {
std::this_thread::sleep_for(std::chrono::milliseconds(50));
}
}
- if (state.peer == nullptr) {
+ if (peer == nullptr) {
LOG_ERROR("[%s] failed to connect to peer %s:%u\n", __func__, host.c_str(), request.port);
return true;
}
- state.peer->get_caps(local_caps);
- if (!state.peer->send_data(local_caps, sizeof(local_caps)) ||
- !state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
- state.peer = nullptr;
+ if (!comm_caps_handshake(peer, /*recv_first =*/ false) ||
+ !peer->send_data(&request.rank, sizeof(request.rank)) || !peer->flush()) {
return true;
}
- state.peer->update_caps(remote_caps);
+ state.peers.push_back(std::move(peer));
}
+
state.rank = request.rank;
state.world = request.world;
- LOG_INFO("[%s] device %u joined pairwise comm as rank %u\n", __func__, request.device, request.rank);
+
+ // Forces the pairwise path, for comparing the two on one build. Setting it on any one rank
+ // is enough, because the negotiation spreads the decision.
+ const bool no_backend_communicator = std::getenv("GGML_RPC_NO_BACKEND_COMMUNICATOR") != nullptr;
+
+ ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id]));
+ auto init_rank = (ggml_backend_comm_init_rank_t)
+ ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_init_rank");
+ auto allreduce_fn = (ggml_backend_comm_allreduce_tensor_t)
+ ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_allreduce_tensor");
+
+ // The ranks must agree on backend communicator or pairwise. Every rank shares whether its
+ // backend provides a communicator, and all ranks use it only if every rank can.
+ const uint8_t self_can_use_backend_comm = !no_backend_communicator && init_rank != nullptr && allreduce_fn != nullptr;
+ std::vector<uint8_t> all_can_use_backend_comm(request.world);
+ if (!comm_bootstrap_exchange(&state, &self_can_use_backend_comm, all_can_use_backend_comm.data(),
+ sizeof(self_can_use_backend_comm))) {
+ state.peers.clear();
+ return true;
+ }
+ const bool use_backend_communicator = std::all_of(all_can_use_backend_comm.begin(), all_can_use_backend_comm.end(),
+ [](uint8_t c) { return c != 0; });
+
+ if (use_backend_communicator) {
+ void * backend_ctx = init_rank(backends[dev_id], (int) request.rank, (int) request.world, comm_bootstrap_exchange, &state);
+ if (backend_ctx == nullptr) {
+ LOG_ERROR("[%s] failed to join communicator as rank %u/%u\n", __func__, request.rank, request.world);
+ state.peers.clear();
+ return true;
+ }
+ state.backend_comm_ctx = backend_ctx;
+ state.allreduce_fn = allreduce_fn;
+ // The sockets have done their job; the collective runs between the backends from here on.
+ state.peers.clear();
+ LOG_INFO("[%s] device %u using backend communicator, rank %u/%u\n", __func__, dev_id, request.rank, request.world);
+ response.ok = 1;
+ return true;
+ }
+
+ // No backend communicator, so fall back to the pairwise exchange, which only covers two ranks.
+ if (request.world != 2) {
+ LOG_ERROR("[%s] world size %u requires a backend communicator\n", __func__, request.world);
+ state.peers.clear();
+ return true;
+ }
+ state.pairwise_backend = backends[dev_id];
+ state.allreduce_fn = comm_allreduce_pairwise;
+ LOG_INFO("[%s] device %u using pairwise communicator, rank %u\n", __func__, dev_id, request.rank);
response.ok = 1;
return true;
}
-bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
- if (request.device >= backends.size()) {
- return false;
- }
- comm_state & state = comm_states[request.device];
- if (state.peer == nullptr) {
- LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
- return false;
+// Bootstrap callback for ggml_backend_comm_init_rank_t, over the peer sockets. Rank 0 collects
+// size bytes from every rank in rank order, then sends the whole array back to every rank.
+// user_data is the comm_state.
+bool rpc_server::comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size) {
+ comm_state & state = *(comm_state *) user_data;
+ uint8_t * recv = (uint8_t *) recv_buf;
+ const size_t total = (size_t) state.world * size;
+
+ if (state.rank == 0) {
+ memcpy(recv, send_buf, size);
+ for (size_t i = 0; i < state.peers.size(); i++) {
+ if (!state.peers[i]->recv_data(recv + (i + 1) * size, size)) {
+ return false;
+ }
+ }
+ for (const socket_ptr & peer : state.peers) {
+ if (!peer->send_data(recv, total) || !peer->flush()) {
+ return false;
+ }
+ }
+ return true;
}
- ggml_backend_t backend = backends[request.device];
+ return state.peers[0]->send_data(send_buf, size) && state.peers[0]->flush() &&
+ state.peers[0]->recv_data(recv, total);
+}
+
+// Reduces this rank's contribution in place, exchanging partials with the peer over peers[0].
+// Used when the backend provides no communicator. comm_ctx_v is the comm_state.
+bool rpc_server::comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors) {
+ comm_state & state = *(comm_state *) comm_ctx_v;
+ ggml_tensor * t_dst = tensors[0];
+ ggml_backend_t backend = state.pairwise_backend;
+
+ const size_t nbytes = ggml_nbytes(t_dst);
+ const int64_t ne = ggml_nelements(t_dst);
size_t ctx_size = 16*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(8, false);
struct ggml_init_params params = {
@@ -2307,16 +2431,7 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
ggml_context_ptr ctx_ptr { ggml_init(params) };
GGML_ASSERT(ctx_ptr != nullptr);
ggml_context * ctx = ctx_ptr.get();
- ggml_tensor * t_dst = deserialize_tensor(ctx, &request.tensor);
- if (t_dst == nullptr || t_dst->buffer == nullptr) {
- LOG_ERROR("[%s] error deserializing tensor\n", __func__);
- return false;
- }
- const size_t nbytes = ggml_nbytes(t_dst);
- const int64_t ne = ggml_nelements(t_dst);
- if (nbytes == 0) {
- return true;
- }
+
// reduce large partials in bf16 to halve the wire bytes; small (decode-sized) ones
// stay f32 since the extra casts and sync cost more than the bytes saved
const bool wire_bf16 = t_dst->type == GGML_TYPE_F32 && ne >= 32768;
@@ -2372,13 +2487,13 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
// rank 0 sends first, rank 1 receives first, so large payloads cannot deadlock
if (state.rank == 0) {
- if (!state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush() ||
- !state.peer->recv_data(state.recv_buf.data(), wire_bytes)) {
+ if (!state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush() ||
+ !state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes)) {
return false;
}
} else {
- if (!state.peer->recv_data(state.recv_buf.data(), wire_bytes) ||
- !state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush()) {
+ if (!state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes) ||
+ !state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush()) {
return false;
}
}
@@ -2408,11 +2523,50 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
return true;
}
+bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
+ if (request.device >= backends.size()) {
+ return false;
+ }
+ comm_state & state = comm_states[request.device];
+ if (state.allreduce_fn == nullptr) {
+ LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
+ return false;
+ }
+
+ struct ggml_init_params params = {
+ /*.mem_size =*/ ggml_tensor_overhead(),
+ /*.mem_buffer =*/ NULL,
+ /*.no_alloc =*/ true,
+ };
+ ggml_context_ptr ctx_ptr { ggml_init(params) };
+ GGML_ASSERT(ctx_ptr != nullptr);
+ ggml_tensor * t_dst = deserialize_tensor(ctx_ptr.get(), &request.tensor);
+ if (t_dst == nullptr || t_dst->buffer == nullptr) {
+ LOG_ERROR("[%s] error deserializing tensor\n", __func__);
+ return false;
+ }
+ if (ggml_nbytes(t_dst) == 0) {
+ return true;
+ }
+
+ // Neither path needs a wait here. The pairwise path waits internally, and the backend orders
+ // its reduction after the work that produced the partial sum.
+ void * comm_ctx = state.backend_comm_ctx != nullptr ? state.backend_comm_ctx : &state;
+ if (!state.allreduce_fn(comm_ctx, &t_dst)) {
+ LOG_ERROR("[%s] allreduce failed for tensor %s\n", __func__, t_dst->name);
+ return false;
+ }
+ return true;
+}
+
bool rpc_server::comm_free(const rpc_msg_comm_free_req & request) {
if (request.device >= backends.size()) {
return false;
}
- comm_states[request.device] = comm_state();
+ // Wait for outstanding work. Destroying a communicator while its kernels are still running
+ // is unsafe.
+ sync_all_backends();
+ comm_free_device(request.device);
return true;
}
@@ -2432,6 +2586,14 @@ bool rpc_server::get_device_memory(const rpc_msg_get_device_memory_req & request
}
rpc_server::~rpc_server() {
+ // Wait for outstanding work. Destroying a communicator while its kernels are still running
+ // is unsafe.
+ sync_all_backends();
+
+ // Free any communicator still open. A client can disconnect without sending COMM_FREE.
+ for (uint32_t dev_id = 0; dev_id < comm_states.size(); dev_id++) {
+ comm_free_device(dev_id);
+ }
for (auto buffer : buffers) {
ggml_backend_buffer_free(buffer);
}
@@ -2997,7 +3159,7 @@ static void ggml_backend_rpc_comm_free(void * comm_ctx_v) {
}
static void * ggml_backend_rpc_comm_init(ggml_backend_t * backends, size_t n_backends) {
- if (n_backends != 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
+ if (n_backends < 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
return nullptr;
}
std::vector<ggml_backend_rpc_comm_context::rank_info> ranks;