diff --git a/common/arg.cpp b/common/arg.cpp index 1066c30e..24db04d6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2059,6 +2059,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, } } )); + add_opt(llama_arg( + {"--dump-folder"}, "FOLDER", + "folder to dump network communication tensors (no dumping if unset)", + [](gpt_params & params, const std::string & value) { + params.dump_folder = 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 61123c5c..af535004 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2011,6 +2011,15 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param } cparams.next_node_ip = new char[params.next_node_ip.length() + 1]; std::strcpy(cparams.next_node_ip, params.next_node_ip.c_str()); + if (cparams.dump_folder != nullptr) { + delete[] cparams.dump_folder; + } + if (!params.dump_folder.empty()) { + cparams.dump_folder = new char[params.dump_folder.length() + 1]; + std::strcpy(const_cast(cparams.dump_folder), params.dump_folder.c_str()); + } else { + cparams.dump_folder = nullptr; + } cparams.n_ctx = params.n_ctx; cparams.n_predict = params.n_predict; diff --git a/common/common.h b/common/common.h index 91994dab..57aa25e5 100644 --- a/common/common.h +++ b/common/common.h @@ -356,6 +356,9 @@ struct gpt_params { // batched-bench params bool batched_bench_output_jsonl = false; + + // tensor dumping + std::string dump_folder = ""; // folder to dump network communication tensors }; // call once at the start of a program if it uses libcommon diff --git a/include/llama.h b/include/llama.h index 764e6729..d75b8577 100644 --- a/include/llama.h +++ b/include/llama.h @@ -378,6 +378,9 @@ extern "C" { // currently works only with CPU execution ggml_abort_callback abort_callback; void * abort_callback_data; + + // Tensor dumping path - if provided, network communication tensors will be dumped + const char * dump_folder; }; // model quantization parameters diff --git a/src/llama.cpp b/src/llama.cpp index fefa2f2b..20bddb25 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -93,6 +93,67 @@ #include #include #include +#include +#include + +int g_llama_send_tensors_counts = 0; +int g_llama_recv_tensors_counts = 0; + +enum class TensorDataType : uint8_t { + FLOAT32 = 0, + // Add other data types as needed +}; + +/** + * @brief Dumps a raw data buffer (tensor) to a binary file with shape and size header. + * The created binary file has the following structure: + * 1. Element Type (uint8_t): A numeric code representing the data type. + * 2. N_Embed (uint64_t): The embedding dimension (width). + * 3. N_Tokens (uint64_t): The number of tokens (height). + * 4. Tensor Size (uint64_t): The total size of the tensor data in bytes. + * 5. Tensor Data (void*): The raw bytes of the tensor. + * + * @param dump_path The full path where the file will be saved. + * @param element_type A code representing the data type of the tensor elements. + * @param n_embed The embedding dimension (width). + * @param n_tokens The number of tokens (height). + * @param tensor_size The total size of the tensor data in bytes. + * @param ptr A const void pointer to the beginning of the tensor data. + */ +void dump_tensors(const std::string& dump_path, + uint8_t element_type, + uint64_t n_embed, + uint64_t n_tokens, + uint64_t tensor_size, + const void* ptr) { + // Open the file for binary writing + // std::ios::binary ensures data is written byte-for-byte without modification. + // std::ios::trunc ensures that if the file exists, it's overwritten. + std::ofstream outfile(dump_path, std::ios::binary | std::ios::trunc); + + if (!outfile.is_open()) { + std::cerr << "Error: Could not open file for writing: " << dump_path << std::endl; + return; + } + + // Write the header and data + try { + outfile.write(reinterpret_cast(&element_type), sizeof(element_type)); + outfile.write(reinterpret_cast(&n_embed), sizeof(n_embed)); + outfile.write(reinterpret_cast(&n_tokens), sizeof(n_tokens)); + outfile.write(reinterpret_cast(&tensor_size), sizeof(tensor_size)); + outfile.write(reinterpret_cast(ptr), tensor_size); + + if (!outfile) { + std::cerr << "Error: A failure occurred while writing to " << dump_path << std::endl; + } else { + std::cout << "Successfully dumped tensor to: " << dump_path << std::endl; + } + + } catch (const std::exception& e) { + std::cerr << "An exception occurred during file write: " << e.what() << std::endl; + } +} std::string get_iso8601_ms_timestamp() { auto now = std::chrono::system_clock::now(); @@ -2641,6 +2702,8 @@ struct llama_cparams { ggml_backend_sched_eval_callback cb_eval; void * cb_eval_user_data; + + const char * dump_folder; }; // TODO: separate into "llama_layer_enc" and "llama_layer_dec" @@ -18035,7 +18098,8 @@ 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) { +static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, struct input_tensors * tensors, const char * dump_folder = nullptr) { + g_llama_send_tensors_counts++; try { std::vector send_msgs; size_t buf_size = 0; @@ -18044,6 +18108,14 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne)); buf_size = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1] * sizeof(float); send_msgs.emplace_back(ubatch->backend_embd, buf_size); + + if (dump_folder && strlen(dump_folder) > 0) { + std::string dump_path = std::string(dump_folder) + "/send_" + std::to_string(g_llama_send_tensors_counts) + ".bin"; + dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), + static_cast(tensors->sub_gf_out->ne[0]), + static_cast(tensors->sub_gf_out->ne[1]), + buf_size, ubatch->backend_embd); + } if (tensors->inp_pos) { send_msgs.emplace_back("inp_pos", strlen("inp_pos")); @@ -18058,7 +18130,8 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba } } -static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, const bool is_out_embd=false) { +static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, const bool is_out_embd=false, const char * dump_folder = nullptr) { + g_llama_recv_tensors_counts++; std::vector recv_msgs; if (!zmq::recv_multipart(socket, std::back_inserter(recv_msgs))) { LLAMA_LOG_INFO("Failed to receive tensor data.\n"); @@ -18075,6 +18148,14 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba size_t buf_size = dims[0] * dims[1] * sizeof(float); float * batch_embd = is_out_embd ? ubatch->out_embd : ubatch->backend_embd; std::memcpy(batch_embd, data_msg.data(), buf_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), + static_cast(dims[0]), + static_cast(dims[1]), + buf_size, data_msg.data()); + } } else if (key == "inp_pos") { int64_t * dims = static_cast(dims_msg.data()); size_t buf_size = dims[0] * sizeof(int32_t); @@ -18502,7 +18583,7 @@ static int llama_decode_internal( // receive data from other nodes if (n_world > 1 && !(my_rank == 0 && i == 0) && !(my_rank == 0 && is_last_l)) { LLAMA_LOG_INFO("[%d][%s][comm][start][recv_tensors][sbatch_tokens: %lu, ubatch_tokens: %u, receive data from other nodes]\n", my_rank, get_iso8601_ms_timestamp().c_str(), lctx.sbatch.n_tokens, ubatch.n_tokens); - llama_recv_tensors(*lctx.recv_socket, &ubatch, is_out_embd); + llama_recv_tensors(*lctx.recv_socket, &ubatch, is_out_embd, lctx.cparams.dump_folder); LLAMA_LOG_INFO("[%d][%s][comm][end][recv_tensors][sbatch_tokens: %lu, ubatch_tokens: %u, receive data from other nodes]\n", my_rank, get_iso8601_ms_timestamp().c_str(), lctx.sbatch.n_tokens, ubatch.n_tokens); } @@ -18578,7 +18659,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); + llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder); 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); } @@ -20331,6 +20412,7 @@ struct llama_context_params llama_context_default_params() { /*.no_perf =*/ true, /*.abort_callback =*/ nullptr, /*.abort_callback_data =*/ nullptr, + /*.dump_folder =*/ nullptr, }; return result; @@ -20841,6 +20923,7 @@ struct llama_context * llama_new_context_with_model( ctx->cparams.n_world = params.n_world; ctx->cparams.rank = params.rank; ctx->cparams.force = params.force; + ctx->cparams.dump_folder = params.dump_folder; ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world; return ctx; }