Skip to content

Commit 7fab617

Browse files
ikawrakowikawrakow
authored
Enable split mode graph for on-the-fly merged up/gate experts (ggml-org#1413)
* Split mode graph for on-the-fly merged ffn_up/gate_exps * Cleanup * Also handle merged bias
1 parent 9b90fd3 commit 7fab617

3 files changed

Lines changed: 86 additions & 11 deletions

File tree

‎ggml/src/ggml-cuda.cu‎

Lines changed: 79 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -854,6 +854,85 @@ GGML_CALL static void ggml_backend_cuda_split_buffer_init_tensor([[maybe_unused]
854854
}
855855

856856
GGML_CALL static void ggml_backend_cuda_split_buffer_set_tensor([[maybe_unused]] ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
857+
if (!tensor->extra && tensor->view_src && tensor->view_src->extra) {
858+
// OK, this is an ugly hack, but I don't really see a way to trick the machine into correctly
859+
// loading non-contiguous merged split tensors.
860+
auto view_src = tensor->view_src;
861+
auto extra = (ggml_split_tensor_t *)view_src->extra;
862+
void * extra_ptr;
863+
memcpy(&extra_ptr, view_src->op_params, sizeof(extra_ptr));
864+
if (extra_ptr) {
865+
std::string merged_name = view_src->name;
866+
if (auto pos = merged_name.find("ffn_gate_up_exps.weight"); pos != std::string::npos) {
867+
std::string name = tensor->name;
868+
auto pos_u = name.find("ffn_up_exps.weight");
869+
auto pos_g = name.find("ffn_gate_exps.weight");
870+
if (pos_u != std::string::npos || pos_g != std::string::npos) {
871+
GGML_ASSERT(extra->split_dim == 1);
872+
auto & ranges = *(const std::vector<std::vector<std::pair<int,int>>> *)extra_ptr;
873+
int ne = 0;
874+
for (int is = 0; is < int(ranges.size()); ++is) {
875+
auto & r = ranges[is];
876+
GGML_ASSERT((extra->splits[is] && !r.empty()) || (!extra->splits[is] && r.empty()));
877+
if (r.empty()) continue;
878+
GGML_ASSERT(r.size() == 2);
879+
auto split = extra->splits[is];
880+
ggml_cuda_set_device(is);
881+
int ir = pos_g != std::string::npos ? 0 : 1;
882+
auto p = r[ir];
883+
size_t offset = 0;
884+
if (ir == 1) {
885+
p.first -= tensor->ne[1];
886+
GGML_ASSERT(p.first >= 0);
887+
offset = split->ne[1]/2 * split->nb[1];
888+
}
889+
for (int i02 = 0; i02 < split->ne[2]; ++i02) {
890+
auto dst = (char *)split->data + i02*split->nb[2] + offset;
891+
auto src = (const char *)data + i02*tensor->nb[2] + ne*tensor->nb[1];
892+
CUDA_CHECK(cudaMemcpyAsync(dst, src, p.second*tensor->nb[1], cudaMemcpyHostToDevice, cudaStreamPerThread));
893+
}
894+
ne += p.second;
895+
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
896+
}
897+
}
898+
return;
899+
}
900+
if (auto pos = merged_name.find("ffn_gate_up_exps.bias"); pos != std::string::npos) {
901+
std::string name = tensor->name;
902+
auto pos_u = name.find("ffn_up_exps.bias");
903+
auto pos_g = name.find("ffn_gate_exps.bias");
904+
if (pos_u != std::string::npos || pos_g != std::string::npos) {
905+
GGML_ASSERT(extra->split_dim == 0);
906+
auto & ranges = *(const std::vector<std::vector<std::pair<int,int>>> *)extra_ptr;
907+
int ne = 0;
908+
for (int is = 0; is < int(ranges.size()); ++is) {
909+
auto & r = ranges[is];
910+
GGML_ASSERT((extra->splits[is] && !r.empty()) || (!extra->splits[is] && r.empty()));
911+
if (r.empty()) continue;
912+
GGML_ASSERT(r.size() == 2);
913+
auto split = extra->splits[is];
914+
ggml_cuda_set_device(is);
915+
int ir = pos_g != std::string::npos ? 0 : 1;
916+
auto p = r[ir];
917+
size_t offset = 0;
918+
if (ir == 1) {
919+
p.first -= tensor->ne[0];
920+
GGML_ASSERT(p.first >= 0);
921+
offset = split->ne[0]/2 * split->nb[0];
922+
}
923+
for (int i01 = 0; i01 < split->ne[1]; ++i01) {
924+
auto dst = (char *)split->data + i01*split->nb[1] + offset;
925+
auto src = (const char *)data + i01*tensor->nb[1] + ne*tensor->nb[0];
926+
CUDA_CHECK(cudaMemcpyAsync(dst, src, p.second*tensor->nb[0], cudaMemcpyHostToDevice, cudaStreamPerThread));
927+
}
928+
ne += p.second;
929+
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
930+
}
931+
}
932+
return;
933+
}
934+
}
935+
}
857936
if (!tensor->extra) return;
858937
static std::map<ggml_type, int> k_map = {
859938
{ GGML_TYPE_Q4_0_R8 , 8},
@@ -886,7 +965,6 @@ GGML_CALL static void ggml_backend_cuda_split_buffer_set_tensor([[maybe_unused]]
886965
{ GGML_TYPE_Q8_KV_R8 , 4},
887966
{ GGML_TYPE_Q8_K_R8 , 8},
888967
};
889-
//printf("%s(%s)\n", __func__, tensor->name);
890968

891969
// split tensors must always be set in their entirety at once
892970
GGML_ASSERT(offset == 0);
@@ -984,7 +1062,6 @@ GGML_CALL static void ggml_backend_cuda_split_buffer_set_tensor([[maybe_unused]]
9841062
auto row_size = ggml_row_size(tensor->type, tensor->ne[0]);
9851063
std::vector<char> host_buffer;
9861064
int ne1 = 0;
987-
int extra_ne1 = 0;
9881065
for (int i = 0; i < extra->n_device; ++i) {
9891066
auto split = extra->splits[i];
9901067
if (!split) continue;

‎src/llama-load-tensors.cpp‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3168,6 +3168,9 @@ bool create_tensors_helper::merge_up_gate_exps(const LLM_TN & tn, int i, int bia
31683168

31693169
layer.ffn_up_gate_exps = ggml_new_tensor_3d(u_ctx, u_meta->type, u_meta->ne[0], u_meta->ne[1] + g_meta->ne[1], u_meta->ne[2]);
31703170
snprintf(layer.ffn_up_gate_exps->name, GGML_MAX_NAME, "blk.%d.ffn_gate_up_exps.weight", i);
3171+
if (u_ctx == ctx_split) {
3172+
split_tensors.insert(layer.ffn_up_gate_exps);
3173+
}
31713174
layer.ffn_gate_exps = ml.create_tensor_as_view(u_ctx, layer.ffn_up_gate_exps, g_name.c_str(),
31723175
{ g_meta->ne[0], g_meta->ne[1], g_meta->ne[2] }, 0);
31733176
layer.ffn_up_exps = ml.create_tensor_as_view(u_ctx, layer.ffn_up_gate_exps, u_name.c_str(),
@@ -3702,13 +3705,6 @@ bool create_tensors_helper::create_tensors() {
37023705
LLAMA_LOG_WARN("========================================================\n\n");
37033706
ml.merge_qkv = false;
37043707
}
3705-
if (ml.merge_up_gate_exps && (model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN)) {
3706-
LLAMA_LOG_WARN("\n========================================================\n");
3707-
LLAMA_LOG_WARN("merge_up_gate_exps is not compatible with split mode 'graph'\n");
3708-
LLAMA_LOG_WARN(" => turning off merge_up_gate_exps\n");
3709-
LLAMA_LOG_WARN("========================================================\n\n");
3710-
ml.merge_up_gate_exps = false;
3711-
}
37123708
switch (model.arch) {
37133709
case LLM_ARCH_LLAMA:
37143710
case LLM_ARCH_REFACT:
@@ -4039,7 +4035,7 @@ bool create_tensors_helper::create_tensors() {
40394035
prepare_split_tensors(1, ctx_split, layer.ffn_up_gate_exps, layer.split_ffn_up_gate_exps, up_gate_split, mem_used);
40404036
prepare_up_gate_split(layer.ffn_up_gate_exps, layer.split_ffn_up_gate_exps);
40414037
if (layer.ffn_up_gate_exps_b) {
4042-
prepare_split_tensors(1, ctx_split, layer.ffn_up_gate_exps_b, layer.split_ffn_up_gate_exps_b, up_gate_split, mem_used);
4038+
prepare_split_tensors(0, ctx_split, layer.ffn_up_gate_exps_b, layer.split_ffn_up_gate_exps_b, up_gate_split, mem_used);
40434039
prepare_up_gate_split(layer.ffn_up_gate_exps_b, layer.split_ffn_up_gate_exps_b);
40444040
}
40454041
} else {

‎src/llama.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5209,7 +5209,9 @@ struct llama_context * llama_init_from_model(
52095209
LLAMA_LOG_INFO("%s: pipeline parallelism enabled (n_copies=%d)\n", __func__, ggml_backend_sched_get_n_copies(ctx->sched));
52105210
}
52115211

5212-
llama_repack_up_gate_exps(*ctx);
5212+
if (ctx->model.split_mode != LLAMA_SPLIT_MODE_GRAPH) {
5213+
llama_repack_up_gate_exps(*ctx);
5214+
}
52135215

52145216
// build worst-case graph
52155217
int n_past = cparams.n_ctx - n_tokens;

0 commit comments

Comments
 (0)