Upload config

config.json

@@ -1,10 +1,6 @@
 {
-  "architectures": [
-    "MambaLMHeadModel"
-  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig",
-    "AutoModelForCausalLM": "modeling_mamba.MambaLMHeadModel"
+    "AutoConfig": "configuration_mamba.MambaConfig"
   },
   "bias": false,
   "conv_bias": true,
@@ -14,10 +10,10 @@
   "d_state": 16,
   "dt_rank": 48,
   "expand": 2,
+  "initializer_range": 0.02,
   "model_type": "mamba",
   "n_layer": 24,
   "pad_vocab_size_multiple": 8,
-  "torch_dtype": "float32",
   "transformers_version": "4.37.2",
   "vocab_size": 50280
 }

configuration_mamba.py

@@ -4,106 +4,45 @@ from typing import Union
 from transformers import PretrainedConfig
 
 
-# Inspired by:
-# - https://huggingface.co/docs/transformers/custom_models#writing-a-custom-configuration
-# - https://huggingface.co/Q-bert/Mamba-130M/blob/9fad7fb5fb9c9416fab4f70ecd62498478be2074/configuration_mamba.py#L5
-# - https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L33
-# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/config_mamba.py#L5
 class MambaConfig(PretrainedConfig):
-    model_type: str = "mamba"
+    model_type = "mamba"
 
     def __init__(
         self,
-        # bias: bool = False,
-        # conv_bias: bool = True,
-        # d_conv: int = 4,
-        # d_model: int = 2560,
-        # d_state: int = 16,
-        # dt_rank: Union[int, str] = "auto",
-        # expand: int = 2,
-        # fused_add_norm: bool = True,
-        # # initializer_range: float = 0.02,
-        # n_layer: int = 64,  # TODO: Rename to num_hidden_layers?
-        # norm_epsilon: float = 1e-5,
-        # pad_vocab_size_multiple: int = 8,
-        # residual_in_fp32: bool = True,
-        # rms_norm: bool = True,
-        # ssm_config: dict = {},
-        # vocab_size: int = 50277,
-
-        d_model: int = 2560,
-        n_layer: int = 64,
-        vocab_size: int = 50277,
-        d_state: int = 16,
-        expand: int = 2,
-        dt_rank: Union[int, str] = 'auto',
-        d_conv: int = 4,
-        pad_vocab_size_multiple: int = 8,
-        conv_bias: bool = True,
-        bias: bool = False,
+        vocab_size=50277,
+        d_state=16,
+        d_model=2560,
+        d_conv=4,
+        expand=2,
+        conv_bias=True,
+        bias=False,
+        n_layer=64,
+        dt_rank: Union[int, str] = "auto",
+        pad_vocab_size_multiple=8,
+        initializer_range=0.02,
         **kwargs,
     ):
-        # self.bias = bias
-        # self.conv_bias = conv_bias
-        # self.d_conv = d_conv
-        # self.d_model = d_model
-        # self.d_state = d_state
-        # self.dt_rank = dt_rank
-        # self.expand = expand
-        # self.fused_add_norm = fused_add_norm
-        # self.n_layer = n_layer
-        # self.norm_epsilon = norm_epsilon
-        # self.pad_vocab_size_multiple = pad_vocab_size_multiple
-        # self.residual_in_fp32 = residual_in_fp32
-        # self.rms_norm = rms_norm
-        # self.ssm_config = ssm_config
-        # self.vocab_size = vocab_size
-
-        # d_model: int
-        # n_layer: int
-        # vocab_size: int
-        # d_state: int = 16
-        # expand: int = 2
-        # dt_rank: Union[int, str] = 'auto'
-        # d_conv: int = 4 
-        # pad_vocab_size_multiple: int = 8
-        # conv_bias: bool = True
-        # bias: bool = False
-
-        self.d_model = d_model
-        self.n_layer = n_layer
         self.vocab_size = vocab_size
-        self.d_state = d_state
+        self.n_layer = n_layer
+        self.conv_bias = conv_bias
         self.expand = expand
-        self.dt_rank = dt_rank
-        self.d_conv = d_conv
         self.pad_vocab_size_multiple = pad_vocab_size_multiple
-        self.conv_bias = conv_bias
-        self.bias = bias
-
+        self.d_conv = d_conv
+        self.d_model = d_model
+        self.d_state = d_state
         self.d_inner = int(self.expand * self.d_model)
+        self.dt_rank = dt_rank
+        self.initializer_range = initializer_range
+        self.bias = bias
 
-        self.d_inner = int(self.expand * self.d_model)
-        
-        if self.dt_rank == 'auto':
+        if self.dt_rank == "auto":
             self.dt_rank = math.ceil(self.d_model / 16)
-            
-        if self.vocab_size % self.pad_vocab_size_multiple != 0:
-            self.vocab_size += (self.pad_vocab_size_multiple
-                                - self.vocab_size % self.pad_vocab_size_multiple)
-
-        # if self.dt_rank == "auto":
-        #     self.dt_rank = math.ceil(self.d_model / 16)  # TODO: 16 is self.d_state?
 
-        # if self.vocab_size % self.pad_vocab_size_multiple != 0:
-        #     self.vocab_size += (
-        #         self.pad_vocab_size_multiple
-        #         - self.vocab_size % self.pad_vocab_size_multiple
-        #     )
-
-        # # TODO: According to https://huggingface.co/docs/transformers/create_a_model#configuration,
-        # # "all NLP models have the hidden_size, num_attention_heads, num_hidden_layers and vocab_size attributes in common."
-        # self.hidden_size = self.d_model
+        if self.vocab_size % self.pad_vocab_size_multiple != 0:
+            self.vocab_size += (
+                self.pad_vocab_size_multiple
+                - self.vocab_size % self.pad_vocab_size_multiple
+            )
 
         super().__init__(
             **kwargs,