electron-rare/prima-cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add configurable network port options

Browse Source
This commit is contained in:
DandinPower
2025-06-10 23:16:57 +08:00
parent fbf853341b
commit 1e7ae71ce5
5 changed files with 70 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/arg.cpp
+35
View File
@@ -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"),
common/common.cpp
+5
View File
@@ -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;
common/common.h
+5
View File
@@ -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
include/llama.h
+5
View File
@@ -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
src/llama.cpp
+20 -9
View File
@@ -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;