diff --git a/src/llama.cpp b/src/llama.cpp index 477038c5..907db25a 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -18171,8 +18171,8 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba free_quantized_array(quantized_array); if (enable_comm_compute_log) { - LLAMA_LOG_INFO("[%d][%s][compute][start][quantize send tensor]\n", my_rank, start_compute_time.c_str()); - LLAMA_LOG_INFO("[%d][%s][compute][end][quantize send tensor]\n", my_rank, end_compute_time.c_str()); + 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) { @@ -18207,24 +18207,24 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba int64_t * dims = static_cast(dims_msg.data()); int64_t * buf_size = static_cast(buffer_size_msg.data()); float * batch_embd = is_out_embd ? ubatch->out_embd : ubatch->backend_embd; - int64_t float_element_size = dims[0] * dims[1] * sizeof(float); + int64_t num_elements = dims[0] * dims[1]; + int64_t float_element_size = num_elements * sizeof(float); - quantized_array_t *quantized_array = (quantized_array_t*)malloc(*buf_size); + quantized_array_t *quantized_array = load_quantized_array_from_buffer(data_msg.data(), *buf_size); if (!quantized_array) { - LLAMA_LOG_INFO("Failed to allocate space for recv data.\n"); + LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n"); return; } - std::memcpy(quantized_array, data_msg.data(), *buf_size); - + std::string start_compute_time = get_iso8601_ms_timestamp(); - dequantize(quantized_array, batch_embd); + 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][quantize send tensor]\n", my_rank, start_compute_time.c_str()); - LLAMA_LOG_INFO("[%d][%s][compute][end][quantize send tensor]\n", my_rank, end_compute_time.c_str()); + 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()); } 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"; diff --git a/src/quantization.cpp b/src/quantization.cpp index de492a30..1217f2ab 100644 --- a/src/quantization.cpp +++ b/src/quantization.cpp @@ -2,8 +2,7 @@ static int64_t _get_q8_0_quantized_array_size(const quantized_array_t *quantized_array) { if (!quantized_array) return 0; - return sizeof(uint8_t) /* quantized_type */ - + 3 * sizeof(uint64_t) /* num_elements, num_blocks, block_size */ + return sizeof(quantized_array_t) /* quantized_type num_elements, num_blocks, block_size */ + quantized_array->num_blocks * sizeof(float) /* scales */ + quantized_array->num_elements * sizeof(int8_t); /* data */ } @@ -31,8 +30,7 @@ int64_t get_quantized_array_size(const quantized_array_t *quantized_array) { } quantized_array_t *allocate_q8_0_array(uint64_t num_elements, - uint64_t block_size) -{ + uint64_t block_size) { if (!num_elements || !block_size) return NULL; uint64_t num_blocks = (num_elements + block_size - 1) / block_size; @@ -57,8 +55,7 @@ quantized_array_t *allocate_q8_0_array(uint64_t num_elements, } quantized_array_t *allocate_q4_0_array(uint64_t num_elements, - uint64_t block_size) -{ + uint64_t block_size) { if (!num_elements || !block_size) return NULL; uint64_t num_blocks = (num_elements + block_size - 1) / block_size; @@ -87,6 +84,25 @@ void free_quantized_array(quantized_array_t *quantized_array) { free(quantized_array); } +quantized_array_t *load_quantized_array_from_buffer(const void *buffer, int64_t buffer_size) { + quantized_array_t *quantized_array = (quantized_array_t*)calloc(1, buffer_size); + if (!quantized_array) return NULL; + + std::memcpy(quantized_array, buffer, buffer_size); + switch (quantized_array->quantized_type) { + case 0: /* q8_0 */ + quantized_array->scales = (float*)(quantized_array + 1); + quantized_array->data = (int8_t*)(quantized_array->scales + quantized_array->num_blocks); + return quantized_array; + case 1: /* q4_0 */ + quantized_array->scales = (float*)(quantized_array + 1); + quantized_array->data = (int8_t*)(quantized_array->scales + quantized_array->num_blocks); + return quantized_array; + default: + return NULL; /* unknown type */ + } +} + static int _quantize_q8_0(const float *float_array, quantized_array_t *quantized_array) { if (!float_array || !quantized_array) return 1; diff --git a/src/quantization.h b/src/quantization.h index e832889d..4efd09b1 100644 --- a/src/quantization.h +++ b/src/quantization.h @@ -1,6 +1,8 @@ #ifndef QUANTIZATION_H #define QUANTIZATION_H +/* To ensure can also compiled with c++ */ +#include #include #include #include @@ -30,6 +32,8 @@ void free_quantized_array(quantized_array_t *quantized_array); int64_t get_quantized_array_size(const quantized_array_t *quantized_array); +quantized_array_t *load_quantized_array_from_buffer(const void *buffer, int64_t buffer_size); + int quantize(const float *float_array, uint64_t num_elements, uint8_t quantized_type,