From 99dd7758d99ddf1c874df010cf439aae1da120de Mon Sep 17 00:00:00 2001 From: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:42:57 -0500 Subject: [PATCH] Add Mellum2 architecture support --- convert_hf_to_gguf.py | 141 +++++++++++++++++++++++++++++++++++ convert_hf_to_gguf_update.py | 1 + gguf-py/gguf/constants.py | 21 ++++++ gguf-py/gguf/gguf_writer.py | 3 + src/CMakeLists.txt | 1 + src/graphs/build_mellum.cpp | 91 ++++++++++++++++++++++ src/llama-arch.cpp | 2 +- src/llama-arch.h | 1 + src/llama-build-context.cpp | 4 + src/llama-build-context.h | 2 + src/llama-hparams.cpp | 25 +++++++ src/llama-load-tensors.cpp | 43 +++++++++++ src/llama-model.cpp | 21 ++++++ src/llama-model.h | 2 +- src/llama-vocab.cpp | 4 + src/llama-vocab.h | 1 + src/llama.cpp | 1 + 17 files changed, 362 insertions(+), 2 deletions(-) create mode 100644 src/graphs/build_mellum.cpp diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 0664e5aab9..ddb7ec8f7e 100644 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -660,6 +660,9 @@ def get_vocab_base_pre(self, tokenizer) -> str: if chkhsh == "f4f37b6c8eb9ea29b3eac6bb8c8487c5ab7885f8d8022e67edc1c68ce8403e95": # ref: https://huggingface.co/MiniMaxAI/MiniMax-M2 res = "minimax-m2" + if chkhsh == "9dcf830ee9990cdbf78cc523a5f7bd9ad8f3f9890c2d3581d2785ad10f07049d": + # ref: https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base + res = "mellum2" if res is None: logger.warning("\n") logger.warning("**************************************************************************************") @@ -2255,6 +2258,144 @@ class Qwen3MoeModel(Qwen2MoeModel): model_arch = gguf.MODEL_ARCH.QWEN3MOE +@Model.register("MellumForCausalLM") +class MellumModel(Model): + model_arch = gguf.MODEL_ARCH.MELLUM + + def set_vocab(self): + tokenizer_path = self.dir_model / "tokenizer.json" + with open(tokenizer_path, "r", encoding="utf-8") as f: + tokenizer_json = json.load(f) + + from tokenizers import Tokenizer + tokenizer = Tokenizer.from_file(str(tokenizer_path)) + + class TokenizerShim: + def encode(self, text: str) -> list[int]: + return tokenizer.encode(text).ids + + vocab: dict[str, int] = tokenizer_json["model"]["vocab"] + vocab_size = self.hparams.get("vocab_size", len(vocab)) + assert max(vocab.values()) < vocab_size + + tokpre = self.get_vocab_base_pre(TokenizerShim()) + reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab.items()} + added_vocab = { + item["content"]: item + for item in tokenizer_json.get("added_tokens", []) + if isinstance(item.get("content"), str) + } + + tokens: list[str] = [] + toktypes: list[int] = [] + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + continue + + token = reverse_vocab[i] + added_token = added_vocab.get(token) + if added_token is not None: + if added_token.get("special", False) or self.does_token_look_special(token): + toktypes.append(gguf.TokenType.CONTROL) + else: + token = token.replace("\u2581", " ") + toktypes.append(gguf.TokenType.USER_DEFINED) + else: + toktypes.append(gguf.TokenType.NORMAL) + tokens.append(token) + + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + if self.hparams.get("num_local_experts") is None and (n_experts := self.hparams.get("num_experts")) is not None: + self.gguf_writer.add_expert_count(n_experts) + + if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None: + self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size) + logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}") + + use_sliding_window = self.hparams.get("use_sliding_window") + sliding_window = self.hparams.get("sliding_window") + if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None: + self.gguf_writer.add_sliding_window(sliding_window) + logger.info(f"gguf: sliding window = {sliding_window}") + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) + logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}") + + rope_parameters = self.hparams.get("rope_parameters", {}) + if full_attention_rope := rope_parameters.get("full_attention"): + if rope_theta := full_attention_rope.get("rope_theta"): + self.gguf_writer.add_rope_freq_base(rope_theta) + logger.info(f"gguf: rope freq base = {rope_theta}") + + if full_attention_rope.get("rope_type") == "yarn": + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + + if factor := full_attention_rope.get("factor"): + self.gguf_writer.add_rope_scaling_factor(factor) + if original_context_length := full_attention_rope.get("original_max_position_embeddings"): + self.gguf_writer.add_rope_scaling_orig_ctx_len(original_context_length) + if attention_factor := full_attention_rope.get("attention_factor"): + self.gguf_writer.add_rope_scaling_yarn_attn_factor(attention_factor) + if beta_fast := full_attention_rope.get("beta_fast"): + self.gguf_writer.add_rope_scaling_yarn_beta_fast(beta_fast) + if beta_slow := full_attention_rope.get("beta_slow"): + self.gguf_writer.add_rope_scaling_yarn_beta_slow(beta_slow) + + if sliding_attention_rope := rope_parameters.get("sliding_attention"): + if rope_theta_swa := sliding_attention_rope.get("rope_theta"): + self.gguf_writer.add_rope_freq_base_swa(rope_theta_swa) + logger.info(f"gguf: rope freq base swa = {rope_theta_swa}") + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if "experts" in name: + n_experts = self.find_hparam(["num_local_experts", "num_experts"]) + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + tensors.append((self.map_tensor_name(merged_name), data_torch)) + return tensors + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + @Model.register("Ernie4_5_ForCausalLM", "Ernie4_5ForCausalLM") class Ernie4_5Model(Model): model_arch = gguf.MODEL_ARCH.ERNIE4_5 diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index 96936717e7..eee5191e3b 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -101,6 +101,7 @@ class TOKENIZER_TYPE(IntEnum): {"name": "kimi-k2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/moonshotai/Kimi-K2-Base", "chkhsh": "81212dc7cdb7e0c1074ca62c5aeab0d43c9f52b8a737be7b12a777c953027890", }, {"name": "grok-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/alvarobartt/grok-2-tokenizer", "chkhsh": "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273"}, {"name": "minimax-m2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/MiniMaxAI/MiniMax-M2", }, + {"name": "mellum2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base", }, ] diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 9219f0e5e4..2cd58b1769 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -115,12 +115,14 @@ class Attention: KV_LORA_RANK = "{arch}.attention.kv_lora_rank" REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count" SLIDING_WINDOW = "{arch}.attention.sliding_window" + SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" OUTPUT_SCALE = "{arch}.attention.output_scale" TEMPERATURE_LENGTH = "{arch}.attention.temperature_length" class Rope: DIMENSION_COUNT = "{arch}.rope.dimension_count" FREQ_BASE = "{arch}.rope.freq_base" + FREQ_BASE_SWA = "{arch}.rope.freq_base_swa" SCALING_TYPE = "{arch}.rope.scaling.type" SCALING_FACTOR = "{arch}.rope.scaling.factor" SCALING_ATTN_FACTOR = "{arch}.rope.scaling.attn_factor" @@ -224,6 +226,7 @@ class MODEL_ARCH(IntEnum): QWEN2MOE = auto() QWEN3 = auto() QWEN3MOE = auto() + MELLUM = auto() PHI2 = auto() PHI3 = auto() PLAMO = auto() @@ -390,6 +393,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN2MOE: "qwen2moe", MODEL_ARCH.QWEN3: "qwen3", MODEL_ARCH.QWEN3MOE: "qwen3moe", + MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.PHI2: "phi2", MODEL_ARCH.PHI3: "phi3", MODEL_ARCH.PLAMO: "plamo", @@ -847,6 +851,23 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, ], + MODEL_ARCH.MELLUM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a051daeeb1..68440f0d66 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -872,6 +872,9 @@ def add_rope_dimension_sections(self, dims: Sequence[int]) -> None: def add_rope_freq_base(self, value: float) -> None: self.add_float32(Keys.Rope.FREQ_BASE.format(arch=self.arch), value) + def add_rope_freq_base_swa(self, value: float) -> None: + self.add_float32(Keys.Rope.FREQ_BASE_SWA.format(arch=self.arch), value) + def add_rope_scaling_type(self, value: RopeScalingType) -> None: self.add_string(Keys.Rope.SCALING_TYPE.format(arch=self.arch), value.value) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 59f9b0a0bf..d5a18c65a3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -83,6 +83,7 @@ add_library(llama graphs/build_qwen.cpp graphs/build_qwen2.cpp graphs/build_qwen3.cpp + graphs/build_mellum.cpp graphs/build_qwen3next.cpp graphs/build_qwen35.cpp graphs/build_phi2.cpp diff --git a/src/graphs/build_mellum.cpp b/src/graphs/build_mellum.cpp new file mode 100644 index 0000000000..8b62ed9ddb --- /dev/null +++ b/src/graphs/build_mellum.cpp @@ -0,0 +1,91 @@ +#include "../llama-build-context.h" +#include "../llama-model.h" +#include "../llama-context.h" + +ggml_cgraph * llm_build_context::build_mellum() { + ggml_cgraph * gf = new_graph_custom(); + + struct ggml_tensor * cur; + struct ggml_tensor * inpL; + + inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + + struct ggml_tensor * inp_pos = build_inp_pos(); + struct ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); + struct ggml_tensor * KQ_mask_swa = build_inp_KQ_mask_swa(); + + for (int il = 0; il < n_layer; ++il) { + const bool is_swa = hparams.swa_layers[il]; + const int64_t n_embd_head = hparams.n_embd_head_v(il); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(il)); + GGML_ASSERT(n_embd_head == hparams.n_rot); + + struct ggml_tensor * inpSA = inpL; + + cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, cb, il); + cb(cur, "attn_norm", il); + + auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, + model.layers[il].wqkv, nullptr, + model.layers[il].wqk, nullptr, + model.layers[il].wq, nullptr, + model.layers[il].wk, nullptr, + model.layers[il].wv, nullptr, + model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0.0f, il); + + const float freq_base_l = is_swa ? hparams.rope_freq_base_train_swa : freq_base; + const float freq_scale_l = is_swa ? 1.0f : freq_scale; + const float ext_factor_l = is_swa ? 0.0f : ext_factor; + const float attn_factor_l = is_swa ? 1.0f : attn_factor; + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast, beta_slow); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = llm_build_kv(ctx0, lctx, kv_self, gf, + model.layers[il].wo, model.layers[il].bo, + Kcur, Vcur, Qcur, is_swa ? KQ_mask_swa : KQ_mask, + n_tokens, kv_head, n_kv, 1.0f/sqrtf(float(n_embd_head)), cb, il, nullptr, is_swa ? hparams.n_swa : 0); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, ffn_inp, + model.layers[il].ffn_gate_inp, nullptr, + model.layers[il].ffn_up_exps, nullptr, + model.layers[il].ffn_gate_exps, nullptr, + model.layers[il].ffn_down_exps, nullptr, + nullptr, + nullptr, nullptr, + nullptr, nullptr, + nullptr, nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, false, 0.0f, + LLM_EXPERT_GATING_FUNC_SOFTMAX, + LLM_FFN_SILU, cb, il, gf, true, + model.layers[il].ffn_up_gate_exps); + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb); + cb(cur, "result_output", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 7e2bb4c441..c0b66cf188 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -32,6 +32,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3VLMOE, "qwen3vlmoe" }, { LLM_ARCH_QWEN35MOE, "qwen35moe" }, { LLM_ARCH_QWEN35, "qwen35" }, + { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_PHI2, "phi2" }, { LLM_ARCH_PHI3, "phi3" }, { LLM_ARCH_PLAMO, "plamo" }, @@ -279,4 +280,3 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { return false; } } - diff --git a/src/llama-arch.h b/src/llama-arch.h index 5a148ad7c3..9ea08b8506 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -31,6 +31,7 @@ enum llm_arch { LLM_ARCH_QWEN3VLMOE, LLM_ARCH_QWEN35MOE, LLM_ARCH_QWEN35, + LLM_ARCH_MELLUM, LLM_ARCH_PHI2, LLM_ARCH_PHI3, LLM_ARCH_PLAMO, diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 8823dccf94..b50e020f2e 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -2300,6 +2300,10 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_qwen3moe(); } break; + case LLM_ARCH_MELLUM: + { + result = llm.build_mellum(); + } break; case LLM_ARCH_QWEN3NEXT: { result = llm.build_qwen3next(); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index c0e2f1e469..d46ff1ba39 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -208,6 +208,8 @@ struct llm_build_context { ggml_cgraph * build_qwen3moe(); + ggml_cgraph * build_mellum(); + ggml_cgraph * build_qwen3vlmoe(); ggml_cgraph * build_qwen3next(); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index f3be1c1148..687df6668d 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -84,6 +84,7 @@ void llm_load_hparams( std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0); std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0); std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0); + std::fill(hparams.swa_layers.begin(), hparams.swa_layers.end(), 0); std::fill(hparams.recurrent_layer_arr.begin(), hparams.recurrent_layer_arr.end(), false); ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, hparams.n_layer, false); @@ -465,6 +466,30 @@ void llm_load_hparams( default: model.type = e_model::MODEL_UNKNOWN; } } break; + case LLM_ARCH_MELLUM: + { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); + + if (hparams.n_swa > 0) { + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer, false)) { + for (uint32_t i = 0; i < hparams.n_layer; ++i) { + hparams.swa_layers[i] = ((i + 1) % 4 != 0); + } + } + + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + } + + switch (hparams.n_layer) { + case 28: model.type = e_model::MODEL_12B_A2_5B; break; + default: model.type = e_model::MODEL_UNKNOWN; + } + } break; case LLM_ARCH_QWEN3NEXT: { ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index b3f1ff06a6..629bec6d89 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -74,6 +74,8 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_qwen3_moe_tensors(const LLM_TN & tn); + bool create_mellum_tensors(const LLM_TN & tn); + bool create_qwen3next_tensors(const LLM_TN & tn); bool create_qwen35moe_tensors(const LLM_TN & tn); @@ -1418,6 +1420,45 @@ bool create_tensors_helper::create_qwen3_moe_tensors(const LLM_TN & tn) { return use_mmap_buffer; } +bool create_tensors_helper::create_mellum_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); + + model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}); + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}); + + for (int i = 0; i < n_layer; ++i) { + ggml_context * ctx_layer = ctx_for_layer(i); + ggml_context * ctx_split = ctx_for_layer_split(i); + + auto & layer = model.layers[i]; + + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}); + + use_mmap_buffer &= !merge_qkv(tn, i, 0); + + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}); + + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}); + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}); + + auto ffn_ctx = model.split_mode == LLAMA_SPLIT_MODE_GRAPH ? ctx_split : ctx_layer; + layer.ffn_norm = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}); + + layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}); + + if (n_expert == 0) { + throw std::runtime_error("n_expert must be > 0 for Mellum"); + } + if (n_expert_used == 0) { + throw std::runtime_error("n_expert_used must be > 0 for Mellum"); + } + + use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, 0, 0, ffn_ctx); + } + return use_mmap_buffer; +} + bool create_tensors_helper::create_qwen3next_tensors(const LLM_TN & tn) { LOADING_PRELUDE model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); @@ -4233,6 +4274,8 @@ bool create_tensors_helper::create_tensors() { case LLM_ARCH_QWEN3MOE: case LLM_ARCH_QWEN3VLMOE: use_mmap_buffer = create_qwen3_moe_tensors(tn); break; + case LLM_ARCH_MELLUM: + use_mmap_buffer = create_mellum_tensors(tn); break; case LLM_ARCH_QWEN3NEXT: use_mmap_buffer = create_qwen3next_tensors(tn); break; case LLM_ARCH_QWEN35MOE: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fef0069d40..f459849f40 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -436,6 +436,26 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" }, }, }, + { + LLM_ARCH_MELLUM, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, + { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, + { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" }, + { LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" }, + { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, + { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, + { LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" }, + }, + }, { LLM_ARCH_QWEN3NEXT, { @@ -1868,6 +1888,7 @@ const char * llama_model_type_name(e_model type) { case MODEL_A13B: return "A13B"; case MODEL_7B_A1B: return "7B.A1B"; case MODEL_8B_A1B: return "8B.A1B"; + case MODEL_12B_A2_5B: return "12B.A2.5B"; case MODEL_16B_A1B: return "16B.A1B"; case MODEL_21B_A3B: return "21B.A3B"; case MODEL_30B_A3B: return "30B.A3B"; diff --git a/src/llama-model.h b/src/llama-model.h index 5ff084feff..9dfff26b13 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -105,6 +105,7 @@ enum e_model { MODEL_A13B, MODEL_7B_A1B, MODEL_8B_A1B, + MODEL_12B_A2_5B, MODEL_16B_A1B, MODEL_21B_A3B, // Ernie MoE small MODEL_30B_A3B, @@ -621,4 +622,3 @@ struct LLM_TN { std::string llama_model_ftype_name(llama_ftype ftype); const char * llama_model_type_name(e_model type); - diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 1ed2a02e5a..1eca29eb5e 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -352,6 +352,7 @@ struct llm_tokenizer_bpe : llm_tokenizer { case LLAMA_VOCAB_PRE_TYPE_CODESHELL: case LLAMA_VOCAB_PRE_TYPE_EXAONE: case LLAMA_VOCAB_PRE_TYPE_MINERVA: + case LLAMA_VOCAB_PRE_TYPE_MELLUM2: regex_exprs = { "\\p{N}", "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)", @@ -1997,6 +1998,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "mellum" || tokenizer_pre == "modern-bert") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2; + } else if ( + tokenizer_pre == "mellum2") { + pre_type = LLAMA_VOCAB_PRE_TYPE_MELLUM2; } else if ( tokenizer_pre == "jais-2") { pre_type = LLAMA_VOCAB_PRE_TYPE_JAIS2; diff --git a/src/llama-vocab.h b/src/llama-vocab.h index 829be07dc1..3dbf8118b8 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -59,6 +59,7 @@ enum llama_vocab_pre_type { LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM = 48, LLAMA_VOCAB_PRE_TYPE_JAIS2 = 49, LLAMA_VOCAB_PRE_TYPE_GEMMA4 = 50, + LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 51, }; struct LLM_KV; diff --git a/src/llama.cpp b/src/llama.cpp index c589b3aa6f..a5d05e608b 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -7382,6 +7382,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_QWEN2MOE: case LLM_ARCH_QWEN3: case LLM_ARCH_QWEN3MOE: + case LLM_ARCH_MELLUM: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_PHI2: case LLM_ARCH_PHI3: