diff --git a/common/arg.cpp b/common/arg.cpp index 47d3c5e6..310d6c73 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -723,6 +723,41 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.next_node_ip = value; } ).set_env("LLAMA_ARG_NEXT_NODE_IP")); + add_opt(llama_arg( + {"--data-port"}, "PORT", + format("data port base for this node (default: %d)", params.data_port), + [](gpt_params & params, int value) { + params.data_port = value; + } + )); + add_opt(llama_arg( + {"--signal-port"}, "PORT", + format("signal port base for this node (default: %d)", params.signal_port), + [](gpt_params & params, int value) { + params.signal_port = value; + } + )); + add_opt(llama_arg( + {"--master-data-port"}, "PORT", + format("master node data port base (default: %d)", params.master_data_port), + [](gpt_params & params, int value) { + params.master_data_port = value; + } + )); + add_opt(llama_arg( + {"--next-node-data-port"}, "PORT", + format("next node data port base (default: %d)", params.next_node_data_port), + [](gpt_params & params, int value) { + params.next_node_data_port = value; + } + )); + add_opt(llama_arg( + {"--next-node-signal-port"}, "PORT", + format("next node signal port base (default: %d)", params.next_node_signal_port), + [](gpt_params & params, int value) { + params.next_node_signal_port = value; + } + )); add_opt(llama_arg( {"--prefetch"}, format("whether to prefetch layer weights (default: %s)", params.prefetch ? "true" : "false"), diff --git a/common/common.cpp b/common/common.cpp index dff98506..ed986303 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1999,6 +1999,11 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param } cparams.master_ip = new char[params.master_ip.length() + 1]; std::strcpy(cparams.master_ip, params.master_ip.c_str()); + cparams.data_port = params.data_port; + cparams.signal_port = params.signal_port; + cparams.master_data_port = params.master_data_port; + cparams.next_node_data_port = params.next_node_data_port; + cparams.next_node_signal_port = params.next_node_signal_port; if (cparams.next_node_ip != nullptr) { delete[] cparams.next_node_ip; diff --git a/common/common.h b/common/common.h index 0a679213..3c777e1e 100644 --- a/common/common.h +++ b/common/common.h @@ -147,6 +147,11 @@ struct gpt_params { uint32_t n_layer_window[32] = {0}; // layer window size on each node std::string master_ip = "localhost"; // ip address of the master node std::string next_node_ip = "localhost"; // ip address of my next node + uint32_t data_port = 9000; // base data port for this node + uint32_t signal_port = 10000; // base signal port for this node + uint32_t master_data_port = 9000; // data port base for master node + uint32_t next_node_data_port = 9000; // data port base for next node + uint32_t next_node_signal_port = 10000; // signal port base for next node bool prefetch = false; // prefetch layer weights bool keep_out_in_metal = true; // whether to keep output weights in metal memory, true by default bool force = false; // force to start prefetching after computation diff --git a/include/llama.h b/include/llama.h index 8bb8ac50..89df1271 100644 --- a/include/llama.h +++ b/include/llama.h @@ -330,6 +330,11 @@ extern "C" { bool keep_out_in_metal; // whether to keep output weights in metal memory char * master_ip; // ip address of the master node char * next_node_ip; // ip address of the next node + uint32_t data_port; // base data port for this node + uint32_t signal_port; // base signal port for this node + uint32_t master_data_port; // master's data port base + uint32_t next_node_data_port; // next node's data port base + uint32_t next_node_signal_port; // next node's signal port base uint32_t n_ctx; // text context, 0 = from model uint32_t n_predict; // number of tokens to predict uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode diff --git a/src/llama.cpp b/src/llama.cpp index 0e615b67..26cbb08f 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3436,6 +3436,9 @@ struct llama_context { // sockets std::string master_ip = "localhost"; std::string next_node_ip = "localhost"; + uint32_t next_node_data_port = 9000; + uint32_t next_node_signal_port = 10000; + uint32_t master_data_port = 9000; uint32_t data_port = 9000; uint32_t signal_port = 10000; zmq::context_t * sock_context = nullptr; @@ -20237,6 +20240,11 @@ struct llama_context_params llama_context_default_params() { /*.keep_out_in_metal =*/ true, /*.master_ip =*/ nullptr, /*.next_node_ip =*/ nullptr, + /*.data_port =*/ 9000, + /*.signal_port =*/ 10000, + /*.master_data_port =*/ 9000, + /*.next_node_data_port =*/ 9000, + /*.next_node_signal_port =*/ 10000, /*.n_ctx =*/ 512, /*.n_predict =*/ 512, /*.n_batch =*/ 2048, @@ -20439,11 +20447,10 @@ void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t m ctx->master_socket = ctx->send_socket; // No need for master_socket in the last rank, reuse send_socket for communication } - const uint32_t next_rank = (my_rank + 1) % n_world; - std::string recv_endp = "tcp://*:" + std::to_string(map_rank_to_port(my_rank, ctx->data_port)); - std::string send_endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->data_port)); - std::string master_endp = "tcp://" + ctx->master_ip + ":" + std::to_string(map_rank_to_port(0, ctx->data_port)); - std::string signal_endp = "tcp://*:" + std::to_string(map_rank_to_port(my_rank, ctx->signal_port)); + std::string recv_endp = "tcp://*:" + std::to_string(ctx->data_port); + std::string send_endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(ctx->next_node_data_port); + std::string master_endp = "tcp://" + ctx->master_ip + ":" + std::to_string(ctx->master_data_port); + std::string signal_endp = "tcp://*:" + std::to_string(ctx->signal_port); try { ctx->recv_socket->bind(recv_endp); @@ -20653,7 +20660,7 @@ int llama_rebuild_topo(llama_context * ctx, uint32_t * n_layer_window, device_in if ((my_rank + 1) % n_world != next_rank) { socket_to_close = ctx->send_socket; ctx->send_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push); - std::string send_endp = "tcp://" + next_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->data_port)); + std::string send_endp = "tcp://" + next_ip + ":" + std::to_string(ctx->next_node_data_port); ctx->send_socket->connect(send_endp); ctx->next_node_ip = next_ip; ctx->cparams.original_next_rank = next_rank; @@ -20723,15 +20730,13 @@ int llama_recv_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window void llama_free_sockets(struct llama_context * ctx, char ** msg) { const uint32_t n_world = ctx->cparams.n_world; const uint32_t my_rank = ctx->cparams.rank; - // to adapt to the new topology, use old next_rank - const uint32_t next_rank = ctx->cparams.original_next_rank; if (n_world == 1) { return; } zmq::socket_t signal_sender(*ctx->sock_context, zmq::socket_type::push); - std::string endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->signal_port)); + std::string endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(ctx->next_node_signal_port); signal_sender.connect(endp); if (my_rank == 0) { @@ -20770,6 +20775,12 @@ struct llama_context * llama_new_context_with_model( ctx->master_ip = params.master_ip; ctx->next_node_ip = params.next_node_ip; + ctx->next_node_data_port = params.next_node_data_port; + ctx->next_node_signal_port = params.next_node_signal_port; + ctx->master_data_port = params.master_data_port; + ctx->data_port = params.data_port; + ctx->signal_port = params.signal_port; + ctx->cparams.n_world = params.n_world; ctx->cparams.rank = params.rank; ctx->cparams.force = params.force;