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

feat: first draft of f32 sparsity support

Browse Source
This commit is contained in:
DandinPower
2025-10-06 10:00:50 +08:00
parent d63cdf13b5
commit 45a78f2623
8 changed files with 244 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+7
View File
@@ -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 \
common/arg.cpp
+8 -1
View File
@@ -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 &params, 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()),
common/common.cpp
+2
View File
@@ -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;
common/common.h
+2
View File
@@ -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
include/llama.h
+2
View File
@@ -399,6 +399,8 @@ extern "C" {
bool enable_comm_compute_log;
const char * comm_datatype;
int comm_sparse_percentage;
};
// model quantization parameters
src/llama.cpp
+62 -12
View File
@@ -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<zmq::message_t> 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<uint8_t>(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;
src/sparsity.cpp
+124
View File
@@ -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;
}
src/sparsity.h
+37
View File
@@ -0,0 +1,37 @@
#ifndef SPARSITY_H
#define SPARSITY_H
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <omp.h>
/**
* @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