loadable, missing cgraph now

Author: ngxson

gguf-py/gguf/tensor_mapping.py

@@ -462,15 +462,15 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.MTP_INP_PROJ: (
-            "model.layers.{bid}.input_proj.weight", # xiaomi mimo
+            "model.layers.{bid}.input_proj", # xiaomi mimo
         ),
 
         MODEL_TENSOR.MTP_TOKEN_NORM: (
-            "model.layers.{bid}.token_layernorm.weight", # xiaomi mimo
+            "model.layers.{bid}.token_layernorm", # xiaomi mimo
         ),
 
         MODEL_TENSOR.MTP_HIDDEN_NORM: (
-            "model.layers.{bid}.hidden_layernorm.weight", # xiaomi mimo
+            "model.layers.{bid}.hidden_layernorm", # xiaomi mimo
         ),
 
         MODEL_TENSOR.SSM_IN: (

src/llama-arch.cpp

@@ -579,6 +579,10 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
             { LLM_TENSOR_FFN_GATE,        "blk.%d.ffn_gate" },
             { LLM_TENSOR_FFN_DOWN,        "blk.%d.ffn_down" },
             { LLM_TENSOR_FFN_UP,          "blk.%d.ffn_up" },
+            { LLM_TENSOR_MTP_INP_PROJ,    "blk.%d.mtp_inp_proj" },
+            { LLM_TENSOR_MTP_TOKEN_NORM,  "blk.%d.mtp_token_norm" },
+            { LLM_TENSOR_MTP_HIDDEN_NORM, "blk.%d.mtp_hidden_norm" },
+            { LLM_TENSOR_LAYER_OUT_NORM,  "blk.%d.layer_output_norm" },
         },
     },
     {
@@ -1678,6 +1682,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_FFN_GATE_EXPS,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
     {LLM_TENSOR_FFN_UP_EXPS,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
     {LLM_TENSOR_FFN_EXP_PROBS_B,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+    {LLM_TENSOR_MTP_INP_PROJ,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_MTP_TOKEN_NORM,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_MTP_HIDDEN_NORM,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     // this tensor is loaded for T5, but never used
     {LLM_TENSOR_DEC_CROSS_ATTN_REL_B,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_NONE}},
     {LLM_TENSOR_CONV1D,                     {LLM_TENSOR_LAYER_INPUT,     GGML_OP_IM2COL}},

src/llama-arch.h

@@ -362,6 +362,9 @@ enum llm_tensor {
     LLM_TENSOR_POS_NET_ATTN_K,
     LLM_TENSOR_POS_NET_ATTN_V,
     LLM_TENSOR_POS_NET_ATTN_OUT,
+    LLM_TENSOR_MTP_INP_PROJ,
+    LLM_TENSOR_MTP_TOKEN_NORM,
+    LLM_TENSOR_MTP_HIDDEN_NORM,
 };
 
 enum llm_tensor_layer {

src/llama-model.cpp

@@ -2364,6 +2364,12 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                         layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
                         layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
                         layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+
+                        // optional MTP (multi token predict), used by Xiaomi Mimo
+                        layer.mtp_inp_proj    = create_tensor(tn(LLM_TENSOR_MTP_INP_PROJ,    "weight", i), {n_embd*2, n_embd}, TENSOR_NOT_REQUIRED);
+                        layer.mtp_token_norm  = create_tensor(tn(LLM_TENSOR_MTP_TOKEN_NORM,  "weight", i), {n_embd}, TENSOR_NOT_REQUIRED);
+                        layer.mtp_hidden_norm = create_tensor(tn(LLM_TENSOR_MTP_HIDDEN_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED);
+                        layer.layer_out_norm  = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM,  "weight", i), {n_embd}, TENSOR_NOT_REQUIRED);
                     }
                 } break;
             case LLM_ARCH_QWEN2MOE:

src/llama-model.h

@@ -312,6 +312,11 @@ struct llama_layer {
     struct ggml_tensor * ffn_up_scale   = nullptr;
     struct ggml_tensor * ffn_down_scale = nullptr;
 
+    // MTP (multi token predict)
+    struct ggml_tensor * mtp_inp_proj    = nullptr;
+    struct ggml_tensor * mtp_token_norm  = nullptr;
+    struct ggml_tensor * mtp_hidden_norm = nullptr;
+
     struct llama_layer_posnet posnet;
 
     struct llama_layer_convnext convnext;