Fix mappings

Author: mitmul

gguf-py/gguf/constants.py

@@ -396,10 +396,15 @@ class MODEL_TENSOR(IntEnum):
     SSM_IN               = auto()
     SSM_CONV1D           = auto()
     SSM_X                = auto()
+    SSM_BCDT             = auto()
     SSM_DT               = auto()
+    SSM_DT_BIAS          = auto()
     SSM_A                = auto()
     SSM_D                = auto()
     SSM_OUT              = auto()
+    SSM_DT_NORM_WEIGHT   = auto()
+    SSM_B_NORM_WEIGHT    = auto()
+    SSM_C_NORM_WEIGHT    = auto()
     TIME_MIX_W0          = auto()
     TIME_MIX_W1          = auto()
     TIME_MIX_W2          = auto()
@@ -677,9 +682,14 @@ class MODEL_TENSOR(IntEnum):
     MODEL_TENSOR.SSM_IN:                    "blk.{bid}.ssm_in",
     MODEL_TENSOR.SSM_CONV1D:                "blk.{bid}.ssm_conv1d",
     MODEL_TENSOR.SSM_X:                     "blk.{bid}.ssm_x",
+    MODEL_TENSOR.SSM_BCDT:                  "blk.{bid}.ssm_bcdt",
     MODEL_TENSOR.SSM_DT:                    "blk.{bid}.ssm_dt",
+    MODEL_TENSOR.SSM_DT_BIAS:               "blk.{bid}.ssm_dt_bias",
     MODEL_TENSOR.SSM_A:                     "blk.{bid}.ssm_a",
     MODEL_TENSOR.SSM_D:                     "blk.{bid}.ssm_d",
+    MODEL_TENSOR.SSM_DT_NORM_WEIGHT:        "blk.{bid}.ssm_dt_norm_weight",
+    MODEL_TENSOR.SSM_B_NORM_WEIGHT:         "blk.{bid}.ssm_b_norm_weight",
+    MODEL_TENSOR.SSM_C_NORM_WEIGHT:         "blk.{bid}.ssm_c_norm_weight",
     MODEL_TENSOR.SSM_OUT:                   "blk.{bid}.ssm_out",
     MODEL_TENSOR.TIME_MIX_W0:               "blk.{bid}.time_mix_w0",
     MODEL_TENSOR.TIME_MIX_W1:               "blk.{bid}.time_mix_w1",
@@ -1280,23 +1290,29 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.OUTPUT,
         MODEL_TENSOR.ROPE_FREQS,
         MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_QKV,
         MODEL_TENSOR.ATTN_Q,
         MODEL_TENSOR.ATTN_K,
-        MODEL_TENSOR.ATTN_V,
         MODEL_TENSOR.ATTN_OUT,
         MODEL_TENSOR.ATTN_ROT_EMBD,
-        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_PRE_NORM,
+        MODEL_TENSOR.FFN_POST_NORM,
         MODEL_TENSOR.FFN_GATE,
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
         # SSM/Mamba tensors for hybrid architecture
         MODEL_TENSOR.SSM_IN,
         MODEL_TENSOR.SSM_CONV1D,
-        MODEL_TENSOR.SSM_X,
+        MODEL_TENSOR.SSM_BCDT,
         MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_DT_BIAS,
         MODEL_TENSOR.SSM_A,
         MODEL_TENSOR.SSM_D,
         MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.SSM_DT_NORM_WEIGHT,
+        MODEL_TENSOR.SSM_B_NORM_WEIGHT,
+        MODEL_TENSOR.SSM_C_NORM_WEIGHT,
     ],
     MODEL_ARCH.GPT2: [
         MODEL_TENSOR.TOKEN_EMBD,

gguf-py/gguf/tensor_mapping.py

@@ -134,6 +134,7 @@ class TensorNameMap:
             "rwkv.blocks.{bid}.ln1",                                # rwkv6
             "model.layers.{bid}.ln1",                               # rwkv7
             "model.layers.{bid}.input_layernorm",                   # llama4
+            "model.layers.{bid}.pre_mixer_norm",                    # plamo2
         ),
 
         # Attention norm 2
@@ -142,6 +143,7 @@ class TensorNameMap:
             "encoder.layer.{bid}.layer_norm_1",             # jina-v2-code
             "rwkv.blocks.{bid}.ln2",                        # rwkv6
             "model.layers.{bid}.ln2",                       # rwkv7
+            "model.layers.{bid}.post_mixer_norm",           # plamo2
         ),
 
         # Attention query-key-value
@@ -160,6 +162,7 @@ class TensorNameMap:
             "model.layers.{bid}.self_attn.qkv_proj",                               # phi3
             "encoder.layers.{bid}.self_attention.query_key_value",                 # chatglm
             "transformer.layers.{bid}.attn.qkv_proj",                              # openelm
+            "model.layers.{bid}.mixer.qkv_proj",                                   # plamo2
         ),
 
         # Attention query
@@ -174,6 +177,7 @@ class TensorNameMap:
             "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok
             "transformer.h.{bid}.attn.attention.q_proj",                 # exaone
             "model.layers.{bid}.self_attn.q_proj",                       # llama4
+            "model.layers.{bid}.mixer.q_weight",                         # plamo2
         ),
 
         # Attention key
@@ -189,6 +193,7 @@ class TensorNameMap:
             "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok
             "transformer.h.{bid}.attn.attention.k_proj",               # exaone
             "model.layers.{bid}.self_attn.k_proj",                     # llama4
+            "model.layers.{bid}.mixer.k_weight",                       # plamo2
         ),
 
         # Attention value
@@ -230,6 +235,7 @@ class TensorNameMap:
             "transformer.layers.{bid}.attn.out_proj",                       # openelm
             "transformer.h.{bid}.attn.attention.out_proj",                  # exaone
             "model.layers.{bid}.self_attn.o_proj",                          # llama4
+            "model.layers.{bid}.mixer.o_proj",                              # plamo2
         ),
 
         # Attention output norm
@@ -271,15 +277,17 @@ class TensorNameMap:
             "model.layers.{bid}.post_attention_layernorm",                   # llama4
         ),
 
-        # Post feed-forward norm
+        # Pre feed-forward norm
         MODEL_TENSOR.FFN_PRE_NORM: (
             "model.layers.{bid}.pre_feedforward_layernorm", # gemma2
+            "model.layers.{bid}.pre_mlp_norm",              # plamo2
         ),
 
         # Post feed-forward norm
         MODEL_TENSOR.FFN_POST_NORM: (
             "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2
-            "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414
+            "model.layers.{bid}.post_mlp_layernorm",         # glm-4-0414
+            "model.layers.{bid}.post_mlp_norm",              # plamo2
         ),
 
         MODEL_TENSOR.FFN_GATE_INP: (
@@ -476,7 +484,10 @@ class TensorNameMap:
         MODEL_TENSOR.SSM_X: (
             "model.layers.{bid}.x_proj",
             "backbone.layers.{bid}.mixer.x_proj",
-            "model.layers.{bid}.mixer.x_proj",  # plamo2
+        ),
+
+        MODEL_TENSOR.SSM_BCDT: (
+            "model.layers.{bid}.mixer.bcdt_proj",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_DT: (
@@ -485,6 +496,10 @@ class TensorNameMap:
             "model.layers.{bid}.mixer.dt_proj",  # plamo2
         ),
 
+        MODEL_TENSOR.SSM_DT_BIAS: (
+            "model.layers.{bid}.mixer.dt_bias",  # plamo2
+        ),
+
         MODEL_TENSOR.SSM_A: (
             "model.layers.{bid}.A_log",
             "backbone.layers.{bid}.mixer.A_log",
@@ -497,6 +512,18 @@ class TensorNameMap:
             "model.layers.{bid}.mixer.D",  # plamo2
         ),
 
+        MODEL_TENSOR.SSM_DT_NORM_WEIGHT: (
+            "model.layers.{bid}.mixer.dt_norm_weight",  # plamo2
+        ),
+
+        MODEL_TENSOR.SSM_B_NORM_WEIGHT: (
+            "model.layers.{bid}.mixer.B_norm_weight",  # plamo2
+        ),
+
+        MODEL_TENSOR.SSM_C_NORM_WEIGHT: (
+            "model.layers.{bid}.mixer.C_norm_weight",  # plamo2
+        ),
+
         MODEL_TENSOR.SSM_OUT: (
             "model.layers.{bid}.out_proj",
             "backbone.layers.{bid}.mixer.out_proj",