diff --git a/common/speculative.cpp b/common/speculative.cpp index 851a47b9a58..70dcb2e41fe 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2542,6 +2542,9 @@ common_speculative_init_result::common_speculative_init_result( model_path = params.speculative.draft.mparams.path; LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str()); + // a draft head can leave out the embeddings and lm head and use the target's + mparams.model_shared = model_tgt; + llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); if (model_dft == NULL) { LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); diff --git a/conversion/bailingmoe3.py b/conversion/bailingmoe3.py index 20bba23e51c..9ba3112ebc5 100644 --- a/conversion/bailingmoe3.py +++ b/conversion/bailingmoe3.py @@ -121,9 +121,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.word_embeddings.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return super().filter_tensors((name, gen)) diff --git a/conversion/base.py b/conversion/base.py index daae28e92ad..565631bc7d5 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -120,6 +120,8 @@ class ModelBase: supports_mtp_export: bool = False mtp_only: bool = False no_mtp: bool = False + # with mtp_only, leave the shared embeddings and lm head to the target model + mtp_shared_embd: bool = False def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False, use_temp_file: bool = False, eager: bool = False, diff --git a/conversion/command_r.py b/conversion/command_r.py index 971f93ebdf1..2b513509d55 100644 --- a/conversion/command_r.py +++ b/conversion/command_r.py @@ -131,9 +131,9 @@ def filter_tensors(cls, item): is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/dots3.py b/conversion/dots3.py index c7ac2319e24..e8d3f350c74 100644 --- a/conversion/dots3.py +++ b/conversion/dots3.py @@ -99,9 +99,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca # --no-mtp: drop the NextN/MTP block; --mtp: keep only that block plus the shared embeddings/norm/lm_head if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/glm.py b/conversion/glm.py index 7544f850cb2..245f01f84bf 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -138,9 +138,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen @@ -292,9 +292,9 @@ def filter_tensors(cls, item): is_mtp = match is not None and int(match.group(1)) >= cls._n_main_layers if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen @@ -352,9 +352,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/ # norm/lm_head (so the resulting GGUF carries just the draft head). - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/qwen.py b/conversion/qwen.py index 419611896fc..ca89d27ba4a 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -338,7 +338,7 @@ def filter_tensors(cls, item): elif len(parts) == 3 and parts[1] in remapper: name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}" elif cls.mtp_only: - keep = name in ( + keep = not cls.mtp_shared_embd and name in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", "embed_tokens.weight", "norm.weight", ) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 78ad26c6563..6e7dddfa661 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -125,6 +125,10 @@ def parse_args() -> argparse.Namespace: "--no-nextn", "--no-mtp", dest="no_mtp", action="store_true", help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.", ) + parser.add_argument( + "--mtp-shared-embd", action="store_true", + help="With --mtp, leave the token embeddings, output norm and LM head out of the draft and take them from the target model at load time. Much smaller draft, but it needs a llama.cpp new enough to read it.", + ) parser.add_argument( "--dspark", action="store_true", help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.", @@ -278,6 +282,12 @@ def main() -> None: if args.mtp: model_class.mtp_only = True + if args.mtp_shared_embd: + if not args.mtp: + logger.error("--mtp-shared-embd only applies together with --mtp") + sys.exit(1) + model_class.mtp_shared_embd = True + model_instance = model_class(dir_model, output_type, fname_out, is_big_endian=args.bigendian, use_temp_file=args.use_temp_file, eager=args.no_lazy, diff --git a/include/llama.h b/include/llama.h index ef7a012c43a..4c2a9fc465a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -340,6 +340,10 @@ extern "C" { // override key-value pairs of the model meta data const struct llama_model_kv_override * kv_overrides; + // already loaded model to take the shared embeddings and lm head from, for a draft + // head that ships neither. must outlive the model being loaded + const struct llama_model * model_shared; + // Keep the booleans together to avoid misalignment during copy-by-value. bool vocab_only; // only load the vocabulary, no weights bool check_tensors; // validate model tensor data diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 7663797ba00..930e1e22764 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1106,6 +1106,70 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten return true; } +// declared in llama-model.h, which this file does not include +const std::vector> & llama_internal_get_tensor_map(const llama_model * model); + +struct ggml_tensor * llama_model_loader::borrow_shared_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne) { + // only the tensors a draft head is allowed to leave out, checked first so no other + // tensor in any model costs a metadata lookup + if (tn.tensor != LLM_TENSOR_TOKEN_EMBD && tn.tensor != LLM_TENSOR_OUTPUT && tn.tensor != LLM_TENSOR_OUTPUT_NORM) { + return nullptr; + } + + // a file that ships the tensor keeps its own copy + const std::string name = tn.str(); + if (get_weight(name.c_str()) != nullptr) { + return nullptr; + } + + // a draft that left out its embeddings left out the whole set, so the missing token_embd is + // what marks one. without this an arch that ties the head to its own token_embd when + // output.weight is absent (qwen3.5, qwen3-next) would borrow the target's head instead + if (get_weight("token_embd.weight") != nullptr) { + return nullptr; + } + + if (model_shared == nullptr) { + throw std::runtime_error(format("%s: missing tensor '%s'; if this is a draft head that shares " + "the target's embeddings, load it as a draft of its target model, not on its own", + __func__, name.c_str())); + } + + ggml_tensor * src = nullptr; + for (const auto & [n, t] : llama_internal_get_tensor_map(model_shared)) { + if (n == name) { + src = t; + break; + } + } + if (src == nullptr) { + throw std::runtime_error(format("%s: draft needs tensor '%s' from the target, which does not have it", + __func__, name.c_str())); + } + + // the draft uses the tensor directly, so the shapes must agree exactly + size_t dim = 0; + for (const int64_t n : ne) { + if (dim >= GGML_MAX_DIMS || src->ne[dim] != n) { + throw std::runtime_error(format("%s: draft and target disagree on '%s': target has %s, draft wants %s", + __func__, name.c_str(), llama_format_tensor_shape(src).c_str(), llama_format_tensor_shape(ne).c_str())); + } + dim++; + } + for (; dim < GGML_MAX_DIMS; dim++) { + if (src->ne[dim] != 1) { + throw std::runtime_error(format("%s: draft and target disagree on '%s': target has %s, draft wants %s", + __func__, name.c_str(), llama_format_tensor_shape(src).c_str(), llama_format_tensor_shape(ne).c_str())); + } + } + + LLAMA_LOG_INFO("%s: tensor %s taken from the target model\n", __func__, name.c_str()); + + // not counted in n_created or size_data: the tensor is not in this file and is neither + // allocated nor freed here + return src; +} + struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { @@ -1326,6 +1390,12 @@ struct ggml_tensor * llama_model_loader::create_tensor( return ret; } + // must run before check_tensor_dims: the tensor is absent from this file by design, and for + // the lm head it must also win over the arch fallback that ties the head to token_embd + if (ggml_tensor * shared = borrow_shared_tensor(tn, ne)) { + return shared; + } + LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str()); const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED), flags & TENSOR_ALLOW_RESHAPE); if (cur == NULL) { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 9e51d0ce750..e211302bd24 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -117,6 +117,10 @@ struct llama_model_loader { std::set tensors; } lazy; + // target model a draft head borrows the shared tensors from, see borrow_shared_tensor() + const struct llama_model * model_shared = nullptr; + + llama_files files; llama_ftype ftype; llama_fver fver; @@ -238,6 +242,11 @@ struct llama_model_loader { const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags); + // a draft head that shares the target's embeddings does not carry its own token_embd, + // output or output_norm; take them from the target model instead. returns null unless the + // file is missing token_embd, so a model that ships its own tensors is never affected + struct ggml_tensor * borrow_shared_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne); + void done_getting_tensors(bool partial = false) const; void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e679b24e87f..22f1f89b9e4 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2692,6 +2692,7 @@ llama_model_params llama_model_default_params() { /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, + /*.model_shared =*/ nullptr, /*.vocab_only =*/ false, /*.check_tensors =*/ false, /*.use_extra_bufts =*/ true, diff --git a/src/llama.cpp b/src/llama.cpp index 633db658c95..7c49b3a2462 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,7 +318,8 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); - ml.lazy.mode = params.lazy_mode; + ml.lazy.mode = params.lazy_mode; + ml.model_shared = params.model_shared; ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params));