From 3164b3aeea0c912300360b6556cb49a2d0a2e5e2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:10:50 +0200 Subject: [PATCH 1/3] llama : split a tied output projection under split mode tensor A model with tied embeddings has no output.weight - the output projection reuses the embedding table under the table's own name. That copy reaches a meta buffer, missed pattern_output_weight and fell through to MIRRORED, so every device held the whole table and ran the whole projection. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a65abfc9ca64..5b765602c8cd 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -416,7 +416,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight"); static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight"); - static const std::regex pattern_output_weight("output\\.weight"); + // a model with tied embeddings has no output.weight - its output projection is a copy of the + // embedding table under the name of the table. Only that copy reaches the meta device, the + // table itself stays on the input device, so the name is unambiguous here. + static const std::regex pattern_output_weight("(output|token_embd)\\.weight"); static const std::regex pattern_output_bias ("output\\.bias"); struct tensor_config { From 2fe08836672b2f5b9742e539327df344f6ed930d Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 00:56:54 +0200 Subject: [PATCH 2/3] tests : cover the tied output tensor split Also shorten the comment that explains the pattern. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 5 ++-- tests/CMakeLists.txt | 7 ++++++ tests/test-llama-archs.cpp | 47 +++++++++++++++++++++++++++++++++++++- 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5b765602c8cd..61e705008102 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -416,9 +416,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight"); static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight"); - // a model with tied embeddings has no output.weight - its output projection is a copy of the - // embedding table under the name of the table. Only that copy reaches the meta device, the - // table itself stays on the input device, so the name is unambiguous here. + // a tied model has no output.weight - its output projection is a copy of token_embd.weight + // only that copy reaches the meta device, the table itself stays on the input device static const std::regex pattern_output_weight("(output|token_embd)\\.weight"); static const std::regex pattern_output_bias ("output\\.bias"); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 04693bb0e6ed..ff16135df170 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -214,6 +214,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) ARGS --test-live-context-workspace ) + llama_test( + test-llama-archs + NAME test-tied-output-split + LABEL main + ARGS --test-tied-output-split + ) + set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3951dbbe8526..b21bc1f723f6 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -13,6 +13,7 @@ #include "../src/llama-context.h" #include "../src/llama-ext.h" #include "../src/llama-model-saver.h" +#include "../src/llama-model.h" #include #include @@ -69,7 +70,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace] [--test-tied-output-split]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -1052,6 +1053,41 @@ static void test_phase_workspace_mismatched_placement(size_t seed) { GGML_ASSERT(llama_contexts_share_workspace(target.get(), draft.get()) == (status == 1)); } +// a tied model has no output.weight - its output projection must get the same split as output.weight +static void test_tied_output_split(size_t seed) { + auto split_state_of = [&](llm_arch arch, bool moe, const char * name, size_t n_devices) { + gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + ggml_backend_dev_t devices[] = { nullptr }; + model_params.devices = devices; + + size_t tensor_seed = seed; + llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tensor_seed, model_params)); + GGML_ASSERT(model); + const ggml_tensor * tensor = model->get_tensor(name); + GGML_ASSERT(tensor != nullptr); + + llama_meta_device_get_split_state_userdata ud = { n_devices, model.get() }; + return llama_meta_device_get_split_state(tensor, &ud); + }; + + const size_t n_devices = 2; + const ggml_backend_meta_split_state ss_tok_embd = split_state_of(LLM_ARCH_LLAMA, false, "token_embd.weight", n_devices); + const ggml_backend_meta_split_state ss_output = split_state_of(LLM_ARCH_LLAMA, false, "output.weight", n_devices); + GGML_ASSERT(ss_tok_embd.axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(ss_tok_embd.axis == ss_output.axis); + GGML_ASSERT(ss_tok_embd.n_segments == ss_output.n_segments); + for (size_t i = 0; i < n_devices; i++) { + GGML_ASSERT(ss_tok_embd.ne[i] == ss_output.ne[i]); + GGML_ASSERT(ss_tok_embd.ne[i] > 0); + } + + // DeepSeek v4 mirrors its output projection, the tied copy must follow + const ggml_backend_meta_split_state ss_dsv4 = split_state_of(LLM_ARCH_DEEPSEEK4, true, "token_embd.weight", n_devices); + GGML_ASSERT(ss_dsv4.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -1541,6 +1577,7 @@ int main(int argc, char ** argv) { std::string out; bool test_phase_workspace = false; bool test_live_context_workspace = false; + bool test_tied_output = false; int verbosity = LOG_LEVEL_ERROR; @@ -1594,6 +1631,10 @@ int main(int argc, char ** argv) { test_live_context_workspace = true; continue; } + if (strcmp(argv[i], "--test-tied-output-split") == 0) { + test_tied_output = true; + continue; + } } printf("%s: using seed %zu\n", __func__, seed); @@ -1611,6 +1652,10 @@ int main(int argc, char ** argv) { test_live_context_workspace_unsupported(seed); return 0; } + if (test_tied_output) { + test_tied_output_split(seed); + return 0; + } if (!out.empty()) { return save_models(arch, seed, verbosity, out); } From 2fca3e081be039a964b014c1fa16c9f54c572134 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 7 Sep 2026 21:33:28 +0200 Subject: [PATCH 3/3] llama : tell the tied output copy from the input table by tensor Both carry the name token_embd.weight, only the copy may be split by vocab row. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 10 +++++---- tests/test-llama-archs.cpp | 44 ++++++++++++++++++++++++++------------ 2 files changed, 36 insertions(+), 18 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 61e705008102..6d356b04679d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -416,9 +416,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight"); static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight"); - // a tied model has no output.weight - its output projection is a copy of token_embd.weight - // only that copy reaches the meta device, the table itself stays on the input device - static const std::regex pattern_output_weight("(output|token_embd)\\.weight"); + static const std::regex pattern_output_weight("output\\.weight"); static const std::regex pattern_output_bias ("output\\.bias"); struct tensor_config { @@ -580,7 +578,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } // output - if (std::regex_match(tensor_name, pattern_output_weight)) { + // a tied model has no output.weight, its output projection is a copy of token_embd.weight. + // the input table carries that name as well, so tell the two apart by the tensor, not by the name + const bool is_output_weight = std::regex_match(tensor_name, pattern_output_weight) || + (tensor == ud->model->output && tensor != ud->model->tok_embd); + if (is_output_weight) { if (is_dsv4) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b21bc1f723f6..10d0885adfcf 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1053,9 +1053,16 @@ static void test_phase_workspace_mismatched_placement(size_t seed) { GGML_ASSERT(llama_contexts_share_workspace(target.get(), draft.get()) == (status == 1)); } -// a tied model has no output.weight - its output projection must get the same split as output.weight +// a tied model has no output.weight - its output projection is a copy of token_embd.weight and must get the same split. +// the input table keeps that name too, it must never be split by vocab row +struct tied_split_states { + ggml_backend_meta_split_state output; // output.weight + ggml_backend_meta_split_state tied; // the same projection, renamed like a tied model + ggml_backend_meta_split_state tok_embd; // the input table +}; + static void test_tied_output_split(size_t seed) { - auto split_state_of = [&](llm_arch arch, bool moe, const char * name, size_t n_devices) { + auto split_states_of = [&](llm_arch arch, bool moe, size_t n_devices) { gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -1065,27 +1072,36 @@ static void test_tied_output_split(size_t seed) { size_t tensor_seed = seed; llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tensor_seed, model_params)); GGML_ASSERT(model); - const ggml_tensor * tensor = model->get_tensor(name); - GGML_ASSERT(tensor != nullptr); + GGML_ASSERT(model->output != nullptr); + GGML_ASSERT(model->tok_embd != nullptr); + GGML_ASSERT(model->output != model->tok_embd); llama_meta_device_get_split_state_userdata ud = { n_devices, model.get() }; - return llama_meta_device_get_split_state(tensor, &ud); + tied_split_states ret; + ret.output = llama_meta_device_get_split_state(model->output, &ud); + // a tied model gives both tensors the same name + ggml_set_name(model->output, "token_embd.weight"); + ret.tied = llama_meta_device_get_split_state(model->output, &ud); + ret.tok_embd = llama_meta_device_get_split_state(model->tok_embd, &ud); + return ret; }; const size_t n_devices = 2; - const ggml_backend_meta_split_state ss_tok_embd = split_state_of(LLM_ARCH_LLAMA, false, "token_embd.weight", n_devices); - const ggml_backend_meta_split_state ss_output = split_state_of(LLM_ARCH_LLAMA, false, "output.weight", n_devices); - GGML_ASSERT(ss_tok_embd.axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(ss_tok_embd.axis == ss_output.axis); - GGML_ASSERT(ss_tok_embd.n_segments == ss_output.n_segments); + const tied_split_states ss_llama = split_states_of(LLM_ARCH_LLAMA, false, n_devices); + GGML_ASSERT(ss_llama.output.axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(ss_llama.tied.axis == ss_llama.output.axis); + GGML_ASSERT(ss_llama.tied.n_segments == ss_llama.output.n_segments); for (size_t i = 0; i < n_devices; i++) { - GGML_ASSERT(ss_tok_embd.ne[i] == ss_output.ne[i]); - GGML_ASSERT(ss_tok_embd.ne[i] > 0); + GGML_ASSERT(ss_llama.tied.ne[i] == ss_llama.output.ne[i]); + GGML_ASSERT(ss_llama.tied.ne[i] > 0); } + GGML_ASSERT(ss_llama.tok_embd.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); // DeepSeek v4 mirrors its output projection, the tied copy must follow - const ggml_backend_meta_split_state ss_dsv4 = split_state_of(LLM_ARCH_DEEPSEEK4, true, "token_embd.weight", n_devices); - GGML_ASSERT(ss_dsv4.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + const tied_split_states ss_dsv4 = split_states_of(LLM_ARCH_DEEPSEEK4, true, n_devices); + GGML_ASSERT(ss_dsv4.output.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(ss_dsv4.tied.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(ss_dsv4.tok_embd.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); } static std::vector get_logits(