vrr/prima.cpp
2
0
Fork 0
mirror of https://github.com/Lizonghang/prima.cpp.git synced 2025-09-08 07:09:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

split input as a separate subgraph

Browse source
This commit is contained in:
Lizonghang 2024-10-30 23:40:13 +04:00
parent 64291ad70c
commit bfc3f9e185
2 changed files with 52 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
ggml/src/ggml.c
View file

@ -4025,8 +4025,6 @@ static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml
ctx->objects_end = obj_new; ctx->objects_end = obj_new;
//printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size);
return obj_new; return obj_new;
} }

109
src/llama.cpp
View file

@ -26,6 +26,7 @@
// TODO: replace with ggml API call // TODO: replace with ggml API call
#define QK_K 256 #define QK_K 256
#define MAX_SCHEDULERS 32
#ifdef __has_include #ifdef __has_include
#if __has_include(<unistd.h>) #if __has_include(<unistd.h>)
@ -3619,9 +3620,9 @@ static bool this_layer_is_mine(
uint32_t layer_id, uint32_t layer_id,
uint32_t n_world, uint32_t n_world,
uint32_t my_rank, uint32_t my_rank,
const uint32_t * n_layer_window) { const uint32_t * n_layer_window) {
uint32_t cumulative_layers = 0; uint32_t cumulative_layers = 0;
uint32_t rank = 0; uint32_t rank = 1;
while (true) { while (true) {
cumulative_layers += n_layer_window[rank]; cumulative_layers += n_layer_window[rank];
if (layer_id < cumulative_layers) { if (layer_id < cumulative_layers) {
@ -3635,8 +3636,7 @@ static int32_t map_layer_to_local_id(
uint32_t layer_id, uint32_t layer_id,
uint32_t n_world, uint32_t n_world,
uint32_t my_rank, uint32_t my_rank,
const uint32_t* n_layer_window) const uint32_t* n_layer_window) {
{
if (!this_layer_is_mine(layer_id, n_world, my_rank, n_layer_window)) { if (!this_layer_is_mine(layer_id, n_world, my_rank, n_layer_window)) {
return -1; return -1;
} }
@ -3650,21 +3650,17 @@ static int32_t map_layer_to_local_id(
uint32_t cumulative_layers = 0; uint32_t cumulative_layers = 0;
uint32_t local_offset = (layer_id / cycle_size) * n_layer_window[my_rank]; uint32_t local_offset = (layer_id / cycle_size) * n_layer_window[my_rank];
for (uint32_t rank = 0; rank < n_world; ++rank) { for (uint32_t rank = 1; rank < n_world; ++rank) {
uint32_t window_size = n_layer_window[rank]; uint32_t window_size = n_layer_window[rank];
if (cycle_offset < cumulative_layers + window_size) { if (cycle_offset < cumulative_layers + window_size) {
if (rank == my_rank) { if (rank == my_rank) {
return cycle_offset - cumulative_layers + local_offset; return cycle_offset - cumulative_layers + local_offset;
} else {
return -1;
} }
} }
cumulative_layers += window_size; cumulative_layers += window_size;
} }
return -1; return cycle_offset - cumulative_layers + local_offset;
} }
// //
@ -7104,18 +7100,12 @@ static bool llm_load_tensors(
} }
} }
// assign the output layer (locate on node 0 only) // assign the input and output layers on CPU by default
if (my_rank == 0) { if (my_rank == 0) {
// there is very little benefit to offloading the input layer, so always keep it on the CPU
model.buft_input = llama_default_buffer_type_cpu(model, true); model.buft_input = llama_default_buffer_type_cpu(model, true);
model.buft_output = llama_default_buffer_type_cpu(model, true);
if (n_gpu_layers > (int)n_layer_window[0]) { LLAMA_LOG_INFO("Layer input assigned to cpu\n");
LLAMA_LOG_INFO("Layer output assigned to gpu\n"); LLAMA_LOG_INFO("Layer output assigned to cpu\n");
model.buft_output = llama_default_buffer_type_offload(model, main_gpu);
} else {
LLAMA_LOG_INFO("Layer output assigned to cpu\n");
model.buft_output = llama_default_buffer_type_cpu(model, true);
}
} }
// count used buffer types // count used buffer types
@ -7132,7 +7122,7 @@ static bool llm_load_tensors(
} }
// create one context per buffer type // create one context per buffer type
size_t ctx_size = ggml_tensor_overhead()*(ml.n_tensors + 1); // +1 for models where tok_embd is duplicated as output size_t ctx_size = ggml_tensor_overhead() * (ml.n_tensors + 1); // +1 for models where tok_embd is duplicated as output
// for moe merged tensors // for moe merged tensors
ctx_size += ggml_tensor_overhead() * my_layers * 3; ctx_size += ggml_tensor_overhead() * my_layers * 3;
@ -8974,7 +8964,7 @@ static bool llm_load_tensors(
} }
} }
ml.init_mappings(false, use_mlock ? &model.mlock_mmaps : nullptr); // disable prefetch here ml.init_mappings(false, use_mlock ? &model.mlock_mmaps : nullptr);
model.mappings.reserve(ml.mappings.size()); model.mappings.reserve(ml.mappings.size());
// create the backend buffers // create the backend buffers
@ -10532,10 +10522,6 @@ struct llm_build_context {
} }
std::vector<ggml_cgraph *> build_llama() { std::vector<ggml_cgraph *> build_llama() {
// create a vector to hold sub-graphs
std::vector<struct ggml_cgraph *> sub_gfs;
struct ggml_cgraph * sub_gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens // mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens; int32_t n_tokens = this->n_tokens;
@ -10543,21 +10529,31 @@ struct llm_build_context {
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot); GGML_ASSERT(n_embd_head == hparams.n_rot);
struct ggml_tensor * cur = nullptr; // create a vector to hold sub-graphs
struct ggml_tensor * inpL = nullptr; std::vector<struct ggml_cgraph *> sub_gfs;
struct ggml_tensor * inpB = nullptr; struct ggml_cgraph * sub_gf = nullptr;
struct ggml_tensor * cur = nullptr;
const uint32_t n_world = this->cparams.n_world; struct ggml_tensor * inpL = nullptr;
const uint32_t my_rank = this->cparams.rank; struct ggml_tensor * inpB = nullptr;
const uint32_t * n_layer_window = this->cparams.n_layer_window; const uint32_t n_world = this->cparams.n_world;
const uint32_t my_rank = this->cparams.rank;
const uint32_t * n_layer_window = this->cparams.n_layer_window;
if (my_rank == 0) { if (my_rank == 0) {
sub_gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
// inp_embd - contains the input embedding // inp_embd - contains the input embedding
inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
// inpB - contains the output embedding from other nodes
inpB = llm_build_backend_embd(ctx0, lctx, hparams, batch, cb); // build the input layer as a seperate subgraph
ggml_build_forward_expand(sub_gf, inpL);
sub_gfs.push_back(sub_gf);
sub_gf = nullptr;
} }
// inpB - contains the output embedding from other nodes
inpB = llm_build_backend_embd(ctx0, lctx, hparams, batch, cb);
// inp_pos - contains the positions // inp_pos - contains the positions
struct ggml_tensor * inp_pos = build_inp_pos(); struct ggml_tensor * inp_pos = build_inp_pos();
@ -10579,10 +10575,7 @@ struct llm_build_context {
continue; continue;
} }
int local_il = map_layer_to_local_id(il, n_world, my_rank, n_layer_window);
if (inpL == nullptr) { if (inpL == nullptr) {
inpB = llm_build_backend_embd(ctx0, lctx, hparams, batch, cb);
inpL = inpB; inpL = inpB;
} }
@ -10592,6 +10585,7 @@ struct llm_build_context {
} }
struct ggml_tensor * inpSA = inpL; // use for shortcut struct ggml_tensor * inpSA = inpL; // use for shortcut
int local_il = map_layer_to_local_id(il, n_world, my_rank, n_layer_window);
// norm // norm
cur = llm_build_norm(ctx0, inpL, hparams, cur = llm_build_norm(ctx0, inpL, hparams,
@ -10649,8 +10643,8 @@ struct llm_build_context {
// skip computing output for unused tokens // skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids(); struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs; n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids); cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
} }
// For Granite architecture // For Granite architecture
@ -16399,8 +16393,8 @@ static struct ggml_cgraph * llama_build_graph_k_shift(llama_context & lctx) {
return result; return result;
} }
static int32_t map_layer_to_subgf_id(uint32_t i, uint32_t my_rank, uint32_t n_world, const uint32_t * n_layer_window) { static int32_t map_layer_to_subgf_id(uint32_t il, uint32_t my_rank, uint32_t n_world, const uint32_t * n_layer_window) {
if (!this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { if (!this_layer_is_mine(il, n_world, my_rank, n_layer_window)) {
return -1; return -1;
} }
@ -16408,7 +16402,7 @@ static int32_t map_layer_to_subgf_id(uint32_t i, uint32_t my_rank, uint32_t n_wo
for (uint32_t rank = 0; rank < n_world; ++rank) { for (uint32_t rank = 0; rank < n_world; ++rank) {
total_window_size += n_layer_window[rank]; total_window_size += n_layer_window[rank];
} }
return i / total_window_size; return my_rank == 0 ? il / total_window_size + 1 : il / total_window_size;
} }
static std::vector<struct ggml_cgraph *> llama_build_graph( static std::vector<struct ggml_cgraph *> llama_build_graph(
@ -16440,7 +16434,7 @@ static std::vector<struct ggml_cgraph *> llama_build_graph(
// norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
// FIXME: fix in ggml_backend_sched // FIXME: fix in ggml_backend_sched
const bool full_offload = lctx.model.n_gpu_layers > (int)n_layer_window[0]; const bool full_offload = false;
if (batch.n_tokens < 32 || full_offload) { if (batch.n_tokens < 32 || full_offload) {
if (il != -1 && strcmp(name, "norm") == 0) { if (il != -1 && strcmp(name, "norm") == 0) {
for (auto * backend : lctx.backends) { for (auto * backend : lctx.backends) {
@ -17489,9 +17483,9 @@ static void manage_graph_tensors(struct ggml_cgraph * cgraph, int advice, bool f
static int llama_decode_internal( static int llama_decode_internal(
llama_context & lctx, llama_context & lctx,
llama_batch batch_all) { // TODO: rename back to batch llama_batch batch_all) { // TODO: rename back to batch
const auto & model = lctx.model; const auto & model = lctx.model;
const auto & hparams = model.hparams; const auto & hparams = model.hparams;
const auto & cparams = lctx.cparams; const auto & cparams = lctx.cparams;
const uint32_t n_world = cparams.n_world; const uint32_t n_world = cparams.n_world;
const uint32_t my_rank = cparams.rank; const uint32_t my_rank = cparams.rank;
@ -17684,7 +17678,7 @@ static int llama_decode_internal(
GGML_ASSERT(my_rank == 0 || n_world > 1); GGML_ASSERT(my_rank == 0 || n_world > 1);
for (size_t i = 0; i < (size_t)gf.size(); ++i) { for (size_t i = 0; i < (size_t)gf.size(); ++i) {
sub_gf = gf[i]; sub_gf = gf[i];
if (n_world > 1 && !(my_rank == 0 && i == 0) && !(my_rank == 0 && is_last_l)) { if (n_world > 1 && !(my_rank == 0 && i == 0) && !(my_rank == 0 && is_last_l)) {
// receive data from previous nodes // receive data from previous nodes
@ -17720,9 +17714,13 @@ static int llama_decode_internal(
break; break;
} }
layer_str = strchr(sub_gf_out->name, '-') + 1; if (strcmp(sub_gf_out->name, "inp_embd") == 0) {
cur_l = std::atoi(layer_str); is_last_l = false;
is_last_l = (cur_l == static_cast<int>(n_layer) - 1); } else {
layer_str = strchr(sub_gf_out->name, '-') + 1;
cur_l = std::atoi(layer_str);
is_last_l = (cur_l == static_cast<int>(n_layer) - 1);
}
// send the result to the next node (or the master) // send the result to the next node (or the master)
if (n_world == 1 || (my_rank == 0 && is_last_l)) { if (n_world == 1 || (my_rank == 0 && is_last_l)) {
@ -17737,7 +17735,7 @@ static int llama_decode_internal(
const size_t buff_size = ubatch.n_tokens * ggml_element_size(tensors.inp_pos); const size_t buff_size = ubatch.n_tokens * ggml_element_size(tensors.inp_pos);
memcpy(tensors.inp_pos->data, ubatch.pos, buff_size); memcpy(tensors.inp_pos->data, ubatch.pos, buff_size);
} }
const bool is_to_master = my_rank !=0 && is_last_l; const bool is_to_master = my_rank != 0 && is_last_l;
zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket; zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket;
llama_send_tensors(*s, &tensors); llama_send_tensors(*s, &tensors);
} }
@ -20141,7 +20139,7 @@ struct llama_context * llama_new_context_with_model(
} }
} }
for (int i = 0; i < 20; ++i) { for (int i = 0; i < MAX_SCHEDULERS; ++i) {
ctx->sched.push_back(ggml_backend_sched_new(ctx->backends.data(), backend_buft.data(), ctx->backends.size(), max_nodes, pipeline_parallel)); ctx->sched.push_back(ggml_backend_sched_new(ctx->backends.data(), backend_buft.data(), ctx->backends.size(), max_nodes, pipeline_parallel));
} }
@ -20152,11 +20150,8 @@ struct llama_context * llama_new_context_with_model(
llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr}; llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
std::vector<ggml_cgraph *> gf = llama_build_graph(*ctx, ubatch, true); std::vector<ggml_cgraph *> gf = llama_build_graph(*ctx, ubatch, true);
GGML_ASSERT(gf.size() <= 20 && "Number of subgraphs exceeds the maximum number of schedulers\n"); GGML_ASSERT(gf.size() <= MAX_SCHEDULERS && "Number of subgraphs exceeds the maximum number of schedulers\n");
ctx->sched.resize(gf.size()); ctx->sched.resize(gf.size());
// prefetch the first subgraph weights
manage_graph_tensors(gf.front(), POSIX_MADV_WILLNEED, false);
// initialize scheduler with the worst-case graph // initialize scheduler with the worst-case graph
bool ok = true; bool ok = true;