Add more mappings for mamba layers in plamo2

Author: mitmul

convert_hf_to_gguf.py

@@ -2191,25 +2191,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
 class Plamo2Model(LlamaModel):
     model_arch = gguf.MODEL_ARCH.PLAMO2
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        
-        # Add custom mappings for Plamo2's unique structure
-        # Plamo2 uses "mixer" for Mamba layers instead of standard attention
-        tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
-        
-        # Add Mamba-specific mappings 
-        for i in range(self.block_count):
-            # SSM/Mamba tensors
-            tensor_map[f"model.layers.{i}.mixer.in_proj"] = f"blk.{i}.ssm_in"
-            tensor_map[f"model.layers.{i}.mixer.conv1d"] = f"blk.{i}.ssm_conv1d"
-            tensor_map[f"model.layers.{i}.mixer.x_proj"] = f"blk.{i}.ssm_x"
-            tensor_map[f"model.layers.{i}.mixer.dt_proj"] = f"blk.{i}.ssm_dt"
-            tensor_map[f"model.layers.{i}.mixer.A_log"] = f"blk.{i}.ssm_a"
-            tensor_map[f"model.layers.{i}.mixer.D"] = f"blk.{i}.ssm_d"
-            tensor_map[f"model.layers.{i}.mixer.out_proj"] = f"blk.{i}.ssm_out"
-            
-        self.tensor_map = tensor_map
 
     def set_vocab(self):
         # Plamo2 uses sentencepiece tokenizer similar to Llama

gguf-py/gguf/tensor_mapping.py

@@ -464,36 +464,43 @@ class TensorNameMap:
         MODEL_TENSOR.SSM_IN: (
             "model.layers.{bid}.in_proj",
             "backbone.layers.{bid}.mixer.in_proj",
+            "model.layers.{bid}.mixer.in_proj",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_CONV1D: (
             "model.layers.{bid}.conv1d",
             "backbone.layers.{bid}.mixer.conv1d",
+            "model.layers.{bid}.mixer.conv1d",  # plamo2
         ),
 
         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_DT: (
             "model.layers.{bid}.dt_proj",
             "backbone.layers.{bid}.mixer.dt_proj",
+            "model.layers.{bid}.mixer.dt_proj",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_A: (
             "model.layers.{bid}.A_log",
             "backbone.layers.{bid}.mixer.A_log",
+            "model.layers.{bid}.mixer.A_log",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_D: (
             "model.layers.{bid}.D",
             "backbone.layers.{bid}.mixer.D",
+            "model.layers.{bid}.mixer.D",  # plamo2
         ),
 
         MODEL_TENSOR.SSM_OUT: (
             "model.layers.{bid}.out_proj",
             "backbone.layers.{bid}.mixer.out_proj",
+            "model.layers.{bid}.mixer.out_proj",  # plamo2
         ),
 
         MODEL_TENSOR.TIME_MIX_W0: (