diff --git a/Makefile b/Makefile index 6d297760..2da48c9a 100644 --- a/Makefile +++ b/Makefile @@ -956,6 +956,7 @@ OBJ_LLAMA = \ src/unicode-data.o \ src/network-utils.o \ src/quantization.o \ + src/sparsity.o \ OBJ_COMMON = \ common/profiler.o \ @@ -1160,6 +1161,7 @@ src/llama.o: \ src/unicode.h \ src/network-utils.h \ src/quantization.h \ + src/sparsity.h \ include/llama.h \ ggml/include/ggml-cuda.h \ ggml/include/ggml-metal.h \ @@ -1180,6 +1182,11 @@ src/quantization.o: \ src/quantization.h $(CXX) $(CXXFLAGS) -c $< -o $@ +src/sparsity.o: \ + src/sparsity.cpp \ + src/sparsity.h + $(CXX) $(CXXFLAGS) -c $< -o $@ -fopenmp + src/llama-grammar.o: \ src/llama-grammar.cpp \ src/llama-grammar.h \ diff --git a/common/arg.cpp b/common/arg.cpp index 832a04d0..73d02e9e 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2120,11 +2120,18 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, )); 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()), + format("Datatype for communication, currently support f32, q8_0, q4_0 or f32_sparsity (default: %s)", params.comm_datatype.c_str()), [](gpt_params & params, const std::string & value) { params.comm_datatype = value; } )); + add_opt(llama_arg( + {"--comm-sparse-percentage"}, "N", + format("Sparse percentage for communication. Currently supported only when the communication data type is f32_sparsity (default: %d). The input value must range from 1 to 100.", params.comm_sparse_percentage), + [](gpt_params ¶ms, int value) { + params.comm_sparse_percentage = 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 dd13c269..9a06c271 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2115,6 +2115,8 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param cparams.comm_datatype = nullptr; } + cparams.comm_sparse_percentage = params.comm_sparse_percentage; + 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 41995013..e5229ac4 100644 --- a/common/common.h +++ b/common/common.h @@ -381,6 +381,8 @@ struct gpt_params { bool enable_comm_compute_log = false; // enable/disable communication and computation logging std::string comm_datatype = "f32"; // data type for communication + + int comm_sparse_percentage = 100; }; // call once at the start of a program if it uses libcommon diff --git a/include/llama.h b/include/llama.h index 67bfe6a7..478f3fa7 100644 --- a/include/llama.h +++ b/include/llama.h @@ -399,6 +399,8 @@ extern "C" { bool enable_comm_compute_log; const char * comm_datatype; + + int comm_sparse_percentage; }; // model quantization parameters diff --git a/src/llama.cpp b/src/llama.cpp index 60124c25..ef29611d 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -14,6 +14,7 @@ #include "network-utils.h" #include "quantization.h" +#include "sparsity.h" #ifdef GGML_USE_RPC # include "ggml-rpc.h" @@ -2718,6 +2719,7 @@ struct llama_cparams { const char * dump_folder; bool enable_comm_compute_log; const char * comm_datatype; + float comm_sparse_percentage; }; // TODO: separate into "llama_layer_enc" and "llama_layer_dec" @@ -18136,7 +18138,7 @@ 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, const char * comm_datatype = nullptr) { +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, int comm_sparse_percentage=100) { g_llama_send_tensors_counts++; try { std::vector send_msgs; @@ -18144,11 +18146,10 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba int64_t float_element_size = num_elements * sizeof(float); 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")); @@ -18157,11 +18158,10 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba 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(); - + quantized_array_t *quantized_array = NULL; if (quantize(ubatch->backend_embd, num_elements, qtype, &quantized_array) || !quantized_array) { LLAMA_LOG_INFO("Failed to allocate space or do quantization\n"); @@ -18177,6 +18177,43 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba sizeof(tensors->sub_gf_out->ne)); send_msgs.emplace_back(quantized_array, buf_size); send_msgs.emplace_back(&buf_size, sizeof(buf_size)); + + free_quantized_array(quantized_array); + if (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()); + } + } else if (comm_datatype_string == "f32_sparsity") { + if (comm_sparse_percentage < 1 && comm_sparse_percentage > 100) { + fprintf(stderr, "Sparse percentage %d should between 1~100\n", comm_sparse_percentage); + return; + } + + float sparse_ratio = comm_sparse_percentage / 100; + sparse_array_t *sparse_array = NULL; + + start_compute_time = get_iso8601_ms_timestamp(); + if (compress(ubatch->backend_embd, tensors->sub_gf_out->ne[0], tensors->sub_gf_out->ne[1], sparse_ratio, &sparse_array)) { + fprintf(stderr, "compress failed for ratio %.2f\n", sparse_ratio); + free_sparse_array(sparse_array); + return; + } + end_compute_time = get_iso8601_ms_timestamp(); + buf_size = get_sparse_array_size(sparse_array); + + send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); + send_msgs.emplace_back("sparse", strlen("sparse")); + send_msgs.emplace_back(tensors->sub_gf_out->ne, + sizeof(tensors->sub_gf_out->ne)); + send_msgs.emplace_back(sparse_array, buf_size); + send_msgs.emplace_back(&buf_size, sizeof(buf_size)); + + free_sparse_array(sparse_array); + if (enable_comm_compute_log) { + LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][sparse_compress]\n", my_rank, start_compute_time.c_str()); + LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][sparse_compress]\n", my_rank, end_compute_time.c_str()); + } + } else { LLAMA_LOG_INFO("Unsupported communication type = %s\n", comm_datatype_string); return; @@ -18195,12 +18232,6 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba zmq::send_multipart(socket, send_msgs); - 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()); - } - 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), @@ -18254,6 +18285,24 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][dequantize]\n", my_rank, end_compute_time.c_str()); } } + else if (comm_type == "sparse") { + sparse_array_t *sparse_array = load_sparse_array_from_buffer(data_msg.data(), *buf_size); + if (!sparse_array) { + LLAMA_LOG_INFO("Failed to load sparse array from buffer.\n"); + return; + } + + std::string start_compute_time = get_iso8601_ms_timestamp(); + decompress(sparse_array, batch_embd) + std::string end_compute_time = get_iso8601_ms_timestamp(); + + free_sparse_array(sparse_array); + + if (enable_comm_compute_log) { + LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][sparse_decompress]\n", my_rank, start_compute_time.c_str()); + LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][sparse_decompress]\n", my_rank, end_compute_time.c_str()); + } + } else { std::memcpy(batch_embd, data_msg.data(), float_element_size); } @@ -18786,7 +18835,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, lctx.cparams.comm_datatype); + llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder, lctx.cparams.enable_comm_compute_log, my_rank, lctx.cparams.comm_datatype, lctx.cparams.comm_sparse_percentage); 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); } @@ -21192,6 +21241,7 @@ struct llama_context * llama_new_context_with_model( 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.comm_sparse_percentage = params.comm_sparse_percentage; ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world; auto &hparams = model->hparams; diff --git a/src/sparsity.cpp b/src/sparsity.cpp new file mode 100644 index 00000000..18b1a515 --- /dev/null +++ b/src/sparsity.cpp @@ -0,0 +1,124 @@ +#include "sparsity.h" + +sparse_array_t *allocate_sparse_array(uint16_t num_tokens, uint16_t num_features, float sparse_ratio) { + if (!num_tokens || !num_features) return NULL; + if (sparse_ratio < 0.0f || sparse_ratio > 1.0f) return NULL; + + float raw_sparse = (float)num_features * sparse_ratio; + uint16_t num_sparse_features = (uint16_t)roundf(raw_sparse); + + // clamp to valid range + if (num_sparse_features > num_features) { + num_sparse_features = num_features; + } else if (num_sparse_features == 0 && sparse_ratio > 0.0f) { + num_sparse_features = 1; // Avoid total sparsity if ratio positive; + } + + uint32_t sparse_elements = (uint32_t)num_tokens * num_sparse_features; + uint64_t total = sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t)); + sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, total); + if (!sparse_array) return NULL; + + /* initialise the header fields */ + sparse_array->num_tokens = num_tokens; + sparse_array->num_features = num_features; + sparse_array->num_sparse_features = num_sparse_features; + sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1); /* just after the header */ + sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements); /* after the sparse_indices */ + + return sparse_array; +} + +void free_sparse_array(sparse_array_t *sparse_array) { + if (!sparse_array) return; + free(sparse_array); +} + +uint64_t get_sparse_array_size(const sparse_array_t *sparse_array) { + if (!sparse_array) return 0; + + uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features; + + return sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t)); +} + +sparse_array_t *load_sparse_array_from_buffer(const void *buffer, uint64_t buffer_size) { + sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, buffer_size); + if (!sparse_array) return NULL; + + memcpy(sparse_array, buffer, buffer_size); + + uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features; + + sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1); + sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements); + + return sparse_array; +} + +typedef struct { + uint16_t index; + float abs_val; +} sort_entry_t; + +static int abs_sort_cmp(const void *a, const void *b) { + const float abs_a = ((const sort_entry_t *)a)->abs_val; + const float abs_b = ((const sort_entry_t *)b)->abs_val; + if (abs_a != abs_b) { + return (abs_a > abs_b) ? -1 : 1; + } + const uint16_t idx_a = ((const sort_entry_t *)a)->index; + const uint16_t idx_b = ((const sort_entry_t *)b)->index; + return (int)idx_a - (int)idx_b; +} + +int compress(const float *float_array, uint16_t num_tokens, uint16_t num_features, float sparse_ratio, sparse_array_t **sparse_array) { + if (!float_array || num_tokens == 0 || num_features == 0 || *sparse_array) return 1; + + /* ---- allocate sparse ------------------------------------------ */ + *sparse_array = allocate_sparse_array(num_tokens, num_features, sparse_ratio); + if (!*sparse_array) return 1; + +#pragma omp parallel for + for (uint16_t cur_token_index = 0; cur_token_index < num_tokens; cur_token_index++) { + sort_entry_t *entries = (sort_entry_t *)malloc(num_features * sizeof(sort_entry_t)); + + uint32_t dense_base = (uint32_t)cur_token_index * num_features; + uint32_t sparse_base = (uint32_t)cur_token_index * (*sparse_array)->num_sparse_features; + + for (uint16_t i = 0; i < num_features; i++) { + entries[i].index = i; + entries[i].abs_val = fabsf(float_array[dense_base + i]); + } + qsort(entries, num_features, sizeof(sort_entry_t), abs_sort_cmp); + + for (uint16_t keep_feature_index = 0; keep_feature_index < (*sparse_array)->num_sparse_features; keep_feature_index++) { + uint16_t orig_index = entries[keep_feature_index].index; + (*sparse_array)->sparse_indices[sparse_base + keep_feature_index] = orig_index; + (*sparse_array)->values[sparse_base + keep_feature_index] = float_array[dense_base + orig_index]; + } + + free(entries); + } + + return 0; +} + +int decompress(const sparse_array_t *sparse_array, float *float_array) { + if (!float_array || !sparse_array) return 1; + + uint32_t num_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_features; + memset(float_array, 0, num_elements * sizeof(float)); + + for (uint16_t cur_token_index = 0; cur_token_index < sparse_array->num_tokens; cur_token_index++) { + uint32_t dense_base = (uint32_t)cur_token_index * sparse_array->num_features; + uint32_t sparse_base = (uint32_t)cur_token_index * sparse_array->num_sparse_features; + + for (uint16_t keep_feature_index = 0; keep_feature_index < sparse_array->num_sparse_features; keep_feature_index++) { + uint16_t original_feature_index = sparse_array->sparse_indices[sparse_base + keep_feature_index]; + float_array[dense_base + original_feature_index] = sparse_array->values[sparse_base + keep_feature_index]; + } + } + + return 0; +} diff --git a/src/sparsity.h b/src/sparsity.h new file mode 100644 index 00000000..3c3946f1 --- /dev/null +++ b/src/sparsity.h @@ -0,0 +1,37 @@ +#ifndef SPARSITY_H +#define SPARSITY_H + +#include +#include +#include +#include +#include + +/** + * @brief Represents a sparse array in zero-based COO format for 2D data with shape [num_tokens, num_features]. + * + * Sparsity is applied along the features dimension. Since each token retains the same number of sparse features, + * token indices are not stored explicitly. The structure holds the selected feature indices and corresponding values + * for all tokens in a flattened manner. + */ +typedef struct { + uint16_t num_tokens; /* Number of tokens (rows in the 2D shape). */ + uint16_t num_features; /* Number of features per token (columns in the 2D shape). */ + uint16_t num_sparse_features; /* Number of retained sparse features per token (must be <= num_features). */ + uint16_t *sparse_indices; /* Flattened array of selected feature indices; length is (num_tokens * num_sparse_features). */ + float *values; /* Flattened array of corresponding sparse values; length is (num_tokens * num_sparse_features). */ +} sparse_array_t; + +sparse_array_t *allocate_sparse_array(uint16_t num_tokens, uint16_t num_features, float sparse_ratio); + +void free_sparse_array(sparse_array_t *sparse_array); + +uint64_t get_sparse_array_size(const sparse_array_t *sparse_array); + +sparse_array_t *load_sparse_array_from_buffer(const void *buffer, uint64_t buffer_size); + +int compress(const float *float_array, uint16_t num_tokens, uint16_t num_features, float sparse_ratio, sparse_array_t **sparse_array); + +int decompress(const sparse_array_t *sparse_array, float *float_array); + +#endif