From 3a97a82f526d7658d2b48b59c8ba9a1d44dec4a7 Mon Sep 17 00:00:00 2001 From: DandinPower Date: Sat, 23 Aug 2025 13:43:14 +0000 Subject: [PATCH] feat: add args to allow user can control the communication datatype --- common/arg.cpp | 7 ++++ common/common.cpp | 10 +++++ common/common.h | 2 + include/llama.h | 4 +- src/llama.cpp | 98 +++++++++++++++++++++++++++++++---------------- 5 files changed, 88 insertions(+), 33 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index ab43ddba..832a04d0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2118,6 +2118,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.enable_comm_compute_log = true; } )); + add_opt(llama_arg( + {"--comm-datatype"}, "TYPE", + format("Datatype for communication, currently support f32, q8_0, q4_0 (default: %s)", params.comm_datatype.c_str()), + [](gpt_params & params, const std::string & value) { + params.comm_datatype = value; + } + )); add_opt(llama_arg( {"--positive-file"}, "FNAME", format("positive prompts file, one prompt per line (default: '%s')", params.cvector_positive_file.c_str()), diff --git a/common/common.cpp b/common/common.cpp index 4d216c1e..dd13c269 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2105,6 +2105,16 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param cparams.enable_comm_compute_log = params.enable_comm_compute_log; + if (cparams.comm_datatype != nullptr) { + delete[] cparams.comm_datatype; + } + if (!params.comm_datatype.empty()) { + cparams.comm_datatype = new char[params.comm_datatype.length() + 1]; + std::strcpy(const_cast(cparams.comm_datatype), params.comm_datatype.c_str()); + } else { + cparams.comm_datatype = nullptr; + } + cparams.n_ctx = params.n_ctx; cparams.n_predict = params.n_predict; cparams.n_seq_max = params.n_parallel; diff --git a/common/common.h b/common/common.h index 6eaeffe0..41995013 100644 --- a/common/common.h +++ b/common/common.h @@ -379,6 +379,8 @@ struct gpt_params { // communication and computation logging bool enable_comm_compute_log = false; // enable/disable communication and computation logging + + std::string comm_datatype = "f32"; // data type for communication }; // call once at the start of a program if it uses libcommon diff --git a/include/llama.h b/include/llama.h index 5795208b..67bfe6a7 100644 --- a/include/llama.h +++ b/include/llama.h @@ -394,9 +394,11 @@ extern "C" { // Tensor dumping path - if provided, network communication tensors will be dumped const char * dump_folder; - + // Enable/disable communication and computation logging for gantt chart analysis bool enable_comm_compute_log; + + const char * comm_datatype; }; // model quantization parameters diff --git a/src/llama.cpp b/src/llama.cpp index 907db25a..60124c25 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -2717,6 +2717,7 @@ struct llama_cparams { const char * dump_folder; bool enable_comm_compute_log; + const char * comm_datatype; }; // TODO: separate into "llama_layer_enc" and "llama_layer_dec" @@ -18135,27 +18136,52 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) { return 0; } -static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, struct input_tensors * tensors, const char * dump_folder = nullptr, const bool enable_comm_compute_log = true, const int my_rank = 0) { +static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, struct input_tensors * tensors, const char * dump_folder = nullptr, const bool enable_comm_compute_log = true, const int my_rank = 0, const char * comm_datatype = nullptr) { g_llama_send_tensors_counts++; try { std::vector send_msgs; int64_t num_elements = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1]; int64_t float_element_size = num_elements * sizeof(float); - quantized_array_t * quantized_array = NULL; - std::string start_compute_time = get_iso8601_ms_timestamp(); - if (quantize(ubatch->backend_embd, num_elements, 0 /*q8_0*/, &quantized_array) || !quantized_array) { - LLAMA_LOG_INFO("Failed to allocate space or doing quantization\n"); + + std::string comm_datatype_string = std::string(comm_datatype); + bool quantized = false; + std::string start_compute_time = ""; + std::string end_compute_time = ""; + int64_t buf_size = 0; + quantized_array_t *quantized_array = NULL; + if (comm_datatype_string == "f32") { + buf_size = float_element_size; + send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); + send_msgs.emplace_back("normal", strlen("normal")); + send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne)); + send_msgs.emplace_back(ubatch->backend_embd, buf_size); + send_msgs.emplace_back(&buf_size, sizeof(int64_t)); + } else if (comm_datatype_string == "q8_0" || comm_datatype_string == "q4_0") { + quantized = true; + int qtype = (comm_datatype_string == "q8_0") ? 0 : 1; + + start_compute_time = get_iso8601_ms_timestamp(); + + if (quantize(ubatch->backend_embd, num_elements, qtype, + &quantized_array) || !quantized_array) { + LLAMA_LOG_INFO("Failed to allocate space or do quantization\n"); + return; + } + + end_compute_time = get_iso8601_ms_timestamp(); + buf_size = get_quantized_array_size(quantized_array); + + send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); + send_msgs.emplace_back("quantized", strlen("quantized")); + send_msgs.emplace_back(tensors->sub_gf_out->ne, + sizeof(tensors->sub_gf_out->ne)); + send_msgs.emplace_back(quantized_array, buf_size); + send_msgs.emplace_back(&buf_size, sizeof(buf_size)); + } else { + LLAMA_LOG_INFO("Unsupported communication type = %s\n", comm_datatype_string); return; } - std::string end_compute_time = get_iso8601_ms_timestamp(); - int64_t buf_size = get_quantized_array_size(quantized_array); - - send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); - send_msgs.emplace_back("quantized", strlen("quantized")); - send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne)); - send_msgs.emplace_back(quantized_array, buf_size); - send_msgs.emplace_back(&buf_size, sizeof(int64_t)); - + if (tensors->inp_pos) { int64_t zero = 0; buf_size = tensors->inp_pos->ne[0] * sizeof(int32_t); @@ -18168,9 +18194,9 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba } zmq::send_multipart(socket, send_msgs); - free_quantized_array(quantized_array); - - if (enable_comm_compute_log) { + + if (quantized) free_quantized_array(quantized_array); + if (quantized && enable_comm_compute_log) { LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][quantize]\n", my_rank, start_compute_time.c_str()); LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][quantize]\n", my_rank, end_compute_time.c_str()); } @@ -18210,22 +18236,28 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba int64_t num_elements = dims[0] * dims[1]; int64_t float_element_size = num_elements * sizeof(float); - quantized_array_t *quantized_array = load_quantized_array_from_buffer(data_msg.data(), *buf_size); - if (!quantized_array) { - LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n"); - return; + if (comm_type == "quantized") { + quantized_array_t *quantized_array = load_quantized_array_from_buffer(data_msg.data(), *buf_size); + if (!quantized_array) { + LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n"); + return; + } + + std::string start_compute_time = get_iso8601_ms_timestamp(); + dequantize(quantized_array, batch_embd); + std::string end_compute_time = get_iso8601_ms_timestamp(); + + free_quantized_array(quantized_array); + + if (enable_comm_compute_log) { + LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][dequantize]\n", my_rank, start_compute_time.c_str()); + LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][dequantize]\n", my_rank, end_compute_time.c_str()); + } } - - std::string start_compute_time = get_iso8601_ms_timestamp(); - dequantize(quantized_array, batch_embd); - std::string end_compute_time = get_iso8601_ms_timestamp(); - - free_quantized_array(quantized_array); - - if (enable_comm_compute_log) { - LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][dequantize]\n", my_rank, start_compute_time.c_str()); - LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][dequantize]\n", my_rank, end_compute_time.c_str()); + else { + std::memcpy(batch_embd, data_msg.data(), float_element_size); } + if (dump_folder && strlen(dump_folder) > 0) { std::string dump_path = std::string(dump_folder) + "/recv_" + std::to_string(g_llama_recv_tensors_counts) + ".bin"; dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), @@ -18754,7 +18786,7 @@ static int llama_decode_internal( struct input_tensors tensors = {sub_gf_out, lctx.inp_pos}; const bool is_to_master = my_rank != 0 && is_last_l; zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket; - llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder, lctx.cparams.enable_comm_compute_log, my_rank); + llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder, lctx.cparams.enable_comm_compute_log, my_rank, lctx.cparams.comm_datatype); if (lctx.cparams.enable_comm_compute_log) { LLAMA_LOG_INFO("[%d][%s][comm][end][send_tensors][sbatch_tokens: %lu, ubatch_tokens: %u, send the result to the next node or the master]\n", my_rank, get_iso8601_ms_timestamp().c_str(), lctx.sbatch.n_tokens, ubatch.n_tokens); } @@ -20514,6 +20546,7 @@ struct llama_context_params llama_context_default_params() { /*.abort_callback_data =*/ nullptr, /*.dump_folder =*/ nullptr, /*.enable_comm_compute_log =*/ false, + /*.comm_datatype =*/ nullptr, }; return result; @@ -21158,6 +21191,7 @@ struct llama_context * llama_new_context_with_model( ctx->cparams.force = params.force; ctx->cparams.dump_folder = params.dump_folder; ctx->cparams.enable_comm_compute_log = params.enable_comm_compute_log; + ctx->cparams.comm_datatype = params.comm_datatype; ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world; auto &hparams = model->hparams;