Upload model

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "MambaModelForCausalLM"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig"
+    "AutoConfig": "configuration_mamba.MambaConfig",
+    "AutoModelForCausalLM": "modeling_mamba.MambaModelForCausalLM"
   },
   "bias": false,
   "conv_bias": true,
@@ -14,6 +18,7 @@
   "model_type": "mamba",
   "n_layer": 24,
   "pad_vocab_size_multiple": 8,
+  "torch_dtype": "float32",
   "transformers_version": "4.37.2",
   "vocab_size": 50280
 }

modeling_mamba.py

@@ -187,7 +187,7 @@ class Mamba(nn.Module):
 
 
 class MambaBlock(nn.Module):
-    def __init__(self, config: MambaConfig):
+    def __init__(self, config: MambaConfig, layer_idx: int = 0):
         """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1]."""
         super().__init__()
         self.config = config
@@ -195,6 +195,9 @@ class MambaBlock(nn.Module):
         self.mixer = Mamba(config)
         self.norm = MambaRMSNorm(config.d_model)
 
+    def forward(self, x):
+        return self.norm(self.mixer(x))
+
 
 class MambaPreTrainedModel(PreTrainedModel):
     config_class = MambaConfig
@@ -226,7 +229,7 @@ class MambaModel(MambaPreTrainedModel):
         self.config = config
 
         self.embedding = nn.Embedding(config.vocab_size, config.d_model)
-        self.layers = nn.ModuleList([MambaBlock(config) for _ in range(config.n_layer)])
+        self.layers = nn.ModuleList([MambaBlock(config, layer_idx) for layer_idx in range(config.n_layer)])
         self.norm_f = MambaRMSNorm(config.d_model)
 
         self.gradient_checkpointing = False