From d3371929bb1b6982cf73f1e54156d3d426cd80a1 Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Sun, 23 Aug 2026 16:19:12 +0530 Subject: [PATCH] [Tensor parallel] Fix meta tensor split state propagation (#27574) * ggml : fix meta tensor split state propagation * Add test-llama-archs to CI --- ci/run.sh | 29 +++++++++++++++ ggml/src/ggml-backend-meta.cpp | 68 ++++++++++++++++++++++++---------- src/llama-model.cpp | 28 ++++++++++++-- tests/test-llama-archs.cpp | 12 +++++- 4 files changed, 113 insertions(+), 24 deletions(-) diff --git a/ci/run.sh b/ci/run.sh index 3d1d75b5b..02ce54916 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -300,6 +300,31 @@ function gg_sum_ctest_release { gg_printf '```\n' } +# test_llama_archs_tensor_split + +function gg_run_test_llama_archs_tensor_split { + cd ${SRC} + + set -e + + GGML_CUDA_DEVICES=1 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + + set +e +} + +function gg_sum_test_llama_archs_tensor_split { + gg_printf '### %s\n\n' "${ci}" + + gg_printf 'Runs test-llama-archs with 1 to 4 CUDA devices\n' + gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)" + gg_printf '```\n' + gg_printf '%s\n' "$(cat $OUT/${ci}.log)" + gg_printf '```\n' +} + # test_scripts function gg_run_test_scripts { @@ -751,6 +776,10 @@ ret=0 test $ret -eq 0 && gg_run ctest_debug test $ret -eq 0 && gg_run ctest_release +if [ ! -z ${GG_BUILD_CUDA} ]; then + test $ret -eq 0 && gg_run test_llama_archs_tensor_split +fi + if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then test $ret -eq 0 && gg_run test_backend_ops_cpu fi diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 7654ea1f3..ded678e68 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -602,27 +602,40 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - GGML_ASSERT(src_ss[0].n_segments == 1); - if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) { - return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1}; - } - int64_t base_ne_in = tensor->src[0]->ne[0]; - for (int dim = 1; dim <= src_ss[0].axis; dim++) { + int64_t base_ne_in = 1; + for (int dim = 0; dim <= src_ss[0].axis; dim++) { base_ne_in *= tensor->src[0]->ne[dim]; } - base_ne_in /= src_ss[0].nr[0]; + if (src_ss[0].n_segments == 1) { + base_ne_in /= src_ss[0].nr[0]; + if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) { + return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1}; + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && tensor->ne[0] == tensor->src[0]->ne[0] && + tensor->ne[1] == 1 && src_ss[0].nr[0] == 1) { + bool complete_rows = true; + for (size_t j = 0; j < n_bufs; j++) { + const int64_t ne = src_ss[0].ne[j]; + complete_rows = complete_rows && (ne == 0 || ne == tensor->src[0]->ne[0]); + } + if (complete_rows) { + // Move a complete dim-0 split to the following singleton dimension. + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + } + } + } + // Reshape outputs use one segment; split-state propagation merges source segments. int64_t base_ne_out = 1; for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim]; - if (base_ne_out_next % base_ne_in == 0) { - return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1}; + base_ne_out *= tensor->ne[dim]; + if (base_ne_out % base_ne_in == 0) { + return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out/base_ne_in)}, 1}; } - if (base_ne_out_next > base_ne_in) { + if (base_ne_out > base_ne_in) { GGML_ASSERT(src_ss[0].n_segments == 1); GGML_ASSERT(src_ss[0].nr[0] == 1); return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } - base_ne_out = base_ne_out_next; } GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); } @@ -792,7 +805,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); - if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { + if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; int64_t ne_sum = 0; for (size_t s = 0; s < ret.n_segments; s++) { @@ -802,6 +815,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } } GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); + } else if (ret.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + GGML_ASSERT(ret.n_segments == 1); + GGML_ASSERT(ret.nr[0] == 1); } return ret; } @@ -1352,15 +1368,29 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg } break; case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { GGML_ASSERT(tensor->type == GGML_TYPE_F32); - const int64_t ne = ggml_nelements(tensor); - std::vector tmp; - tmp.reserve(ne); - for (int64_t i = 0; i < ne; i++) { - tmp.push_back(((const float *) data)[i] / n_bufs); + GGML_ASSERT(offset % sizeof(float) == 0); + GGML_ASSERT(size % sizeof(float) == 0); + const size_t n_values = size / sizeof(float); + size_t n_contributors = 0; + for (size_t j = 0; j < n_bufs; j++) { + n_contributors += split_state.ne[j] != 0; + } + const bool has_contributor_mask = n_contributors != 0; + if (!has_contributor_mask) { + n_contributors = n_bufs; + } + std::vector tmp(n_values); + for (size_t i = 0; i < n_values; i++) { + tmp[i] = ((const float *) data)[i] / n_contributors; + } + std::vector zero; + if (has_contributor_mask) { + zero.resize(n_values, 0.0f); } for (size_t j = 0; j < n_bufs; j++) { ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size); + const float * partial = has_contributor_mask && split_state.ne[j] == 0 ? zero.data() : tmp.data(); + ggml_backend_tensor_set(simple_tensor, partial, offset, size); } } break; default: { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index de0d3c1a6..33f5661b2 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -520,7 +520,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); } if (std::regex_match(tensor_name, pattern_ffn_down_exps_bias)) { - return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL); + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL, "ffn_down_exps.weight"); } // output @@ -554,6 +554,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim); return {{key_dim, 2}, {value_dim, 1}}; } + if (std::regex_match(tensor_name, pattern_r_cache)) { + return {{key_dim * (hparams.ssm_d_conv - 1), 2}, {value_dim * (hparams.ssm_d_conv - 1), 1}}; + } } else { const int64_t head_ratio = n_v_heads / n_k_heads; if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) { @@ -642,12 +645,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str blck_size_perf *= 2; } + const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf); + const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head if (std::regex_match(tensor_name, pattern_attn_sinks)) { GGML_ASSERT(segments.size() == 1); - return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa}; + return {granularity_head}; } - - const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf); if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: @@ -660,6 +663,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str GGML_ASSERT(segments.size() == 1); return {granularity_q}; } + if (std::regex_match(tensor_name, pattern_attn_gate_weight)) { + GGML_ASSERT(segments.size() == 1); + if (tensor->ne[1] == hparams.n_head(il)) { + return {granularity_head}; + } + return {granularity_q}; + } const int64_t granularity_kv = granularity_q / n_gqa; if (std::regex_match(tensor_name, pattern_kv_weight) || @@ -728,6 +738,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str memset(split_state.ne, 0, sizeof(split_state.ne)); split_state.nr[0] = 1; split_state.n_segments = 1; + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + GGML_ASSERT(tc.tensor_axis_0 != tensor); + const ggml_backend_meta_split_state source_split_state = llama_meta_device_get_split_state(tc.tensor_axis_0, userdata); + GGML_ASSERT(source_split_state.axis >= 0 && source_split_state.axis < GGML_MAX_DIMS); + for (size_t j = 0; j < ud->n_devices; j++) { + for (size_t is = 0; is < source_split_state.n_segments; is++) { + split_state.ne[j] += source_split_state.ne[is*ud->n_devices + j] * source_split_state.nr[is]; + } + } + } } return split_state; GGML_UNUSED(userdata); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 07e3a7a11..18676f2be 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -101,6 +101,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_head = 1; n_ff = 96; n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded + } else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) { + n_embd = 160; // exercise per-head tensor split granularity with head size 80 + } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { + n_head = 4; } else if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA @@ -120,6 +124,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_vocab = 4096; // must be >= the hard-coded codec head size (3072) } + uint32_t n_head_kv = n_head; + if (arch == LLM_ARCH_QWEN3) { + n_head_kv = 1; // MQA coverage + } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { + n_head_kv = 2; // GQA coverage + } const uint32_t n_embd_head = n_embd / n_head; ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE, llm_arch_name(arch)); @@ -160,7 +170,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); - ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_kv); } ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);