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

@@ -76,8 +76,8 @@ class Mamba(nn.Module):
         """
 
         (b, l, d) = x.shape
-        x_copy = x  # There was a separate class for residual, I deleted that part and added it here.
-        x = self.norm(x)
+        # x_copy = x  # There was a separate class for residual, I deleted that part and added it here.
+        # x = self.norm(x)
         x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)
         (x, res) = x_and_res.split(
             split_size=[self.config.d_inner, self.config.d_inner], dim=-1
@@ -93,7 +93,8 @@ class Mamba(nn.Module):
 
         y = y * F.silu(res)
 
-        output = self.out_proj(y) + x_copy
+        # output = self.out_proj(y) + x_copy
+        output = self.out_proj(y)
 
         return output
 
@@ -196,7 +197,7 @@ class MambaBlock(nn.Module):
         self.norm = MambaRMSNorm(config.d_model)
 
     def forward(self, x):
-        return self.norm(self.mixer(x))
+        return self.mixer(self.norm(x)) + x
 
 
 class MambaPreTrainedModel(PreTrainedModel):