Upload model

config.json

@@ -0,0 +1,21 @@
+{
+  "_name_or_path": "huseinzol05-dummy-mamba-1.4b/checkpoint-400",
+  "architectures": [
+    "MambaLMHeadModel"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling.MambaLMHeadModel"
+  },
+  "d_model": 2048,
+  "fused_add_norm": true,
+  "hidden_size": 2048,
+  "model_type": "mamba",
+  "n_layer": 48,
+  "pad_vocab_size_multiple": 8,
+  "residual_in_fp32": true,
+  "rms_norm": true,
+  "ssm_cfg": {},
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.0.dev0",
+  "vocab_size": 32000
+}

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2acae916870f1b5fc2b2a3ec2e8efb4a90b7d6c7c24ebd979a709f5e4ae0988
+size 4988506824

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0f4988510bc7fca2780f682617c3422a6c19f0ae2e46c1ce2058e845703106a
+size 350916592

model.safetensors.index.json

@@ -0,0 +1,538 @@
+{
+  "metadata": {
+    "total_size": 5339365376
+  },
+  "weight_map": {
+    "backbone.embedding.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.0.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.1.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.10.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.11.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.12.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.13.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.14.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.15.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.16.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.17.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.18.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.19.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.2.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.20.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.21.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.22.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.23.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.24.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.25.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.26.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.27.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.28.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.29.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.3.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.30.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.31.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.32.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.33.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.34.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.35.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.36.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.37.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.38.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.39.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.4.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.40.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.41.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.42.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.43.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.44.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.44.norm.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.44.norm.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.A_log": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.D": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.conv1d.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.conv1d.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.dt_proj.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.dt_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.in_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.mixer.x_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.norm.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.45.norm.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.A_log": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.D": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.conv1d.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.conv1d.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.dt_proj.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.dt_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.in_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.mixer.x_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.norm.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.46.norm.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.A_log": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.D": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.conv1d.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.conv1d.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.dt_proj.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.dt_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.in_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.mixer.x_proj.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.norm.bias": "model-00002-of-00002.safetensors",
+    "backbone.layers.47.norm.weight": "model-00002-of-00002.safetensors",
+    "backbone.layers.5.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.5.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.6.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.7.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.8.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.A_log": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.D": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.conv1d.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.conv1d.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.dt_proj.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.dt_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.in_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.mixer.x_proj.weight": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.norm.bias": "model-00001-of-00002.safetensors",
+    "backbone.layers.9.norm.weight": "model-00001-of-00002.safetensors",
+    "backbone.norm_f.bias": "model-00002-of-00002.safetensors",
+    "backbone.norm_f.weight": "model-00002-of-00002.safetensors"
+  }
+}

modeling.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2023, Albert Gu, Tri Dao.
+
+import math
+from functools import partial
+
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+from transformers import PretrainedConfig, PreTrainedModel
+
+from mamba_ssm.modules.mamba_simple import Mamba, Block
+from mamba_ssm.utils.generation import GenerationMixin
+from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+
+class MambaConfig(PretrainedConfig):
+    model_type = 'mamba'
+
+
+def create_block(
+    d_model,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=False,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+):
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    factory_kwargs = {"device": device, "dtype": dtype}
+    mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+    norm_cls = partial(
+        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
+    )
+    block = Block(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM):
+        # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+class MixerModel(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_layer: int,
+        vocab_size: int,
+        ssm_cfg=None,
+        norm_epsilon: float = 1e-5,
+        rms_norm: bool = False,
+        initializer_cfg=None,
+        fused_add_norm=False,
+        residual_in_fp32=False,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)
+
+        # We change the order of residual and layer norm:
+        # Instead of LN -> Attn / MLP -> Add, we do:
+        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
+        # the main branch (output of MLP / Mixer). The model definition is unchanged.
+        # This is for performance reason: we can fuse add + layer_norm.
+        self.fused_add_norm = fused_add_norm
+        if self.fused_add_norm:
+            if layer_norm_fn is None or rms_norm_fn is None:
+                raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")
+
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    d_model,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    **factory_kwargs,
+                )
+                for i in range(n_layer)
+            ]
+        )
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            d_model, eps=norm_epsilon, **factory_kwargs
+        )
+
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(self, input_ids, inference_params=None):
+        hidden_states = self.embedding(input_ids)
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+        if not self.fused_add_norm:
+            residual = (hidden_states + residual) if residual is not None else hidden_states
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm=False here since we don't need the residual
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            hidden_states = fused_add_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        return hidden_states
+
+
+class MambaLMHeadModel(PreTrainedModel, GenerationMixin):
+    config_class = MambaConfig
+
+    def __init__(
+        self,
+        config,
+        initializer_cfg=None,
+        pad_vocab_size_multiple: int = 1,
+        device=None,
+        dtype=None,
+        **backbone_kwargs,
+    ) -> None:
+        super().__init__(config)
+        d_model = config.d_model
+        n_layer = config.n_layer
+        vocab_size = config.vocab_size
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        if vocab_size % pad_vocab_size_multiple != 0:
+            vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)
+        self.backbone = MixerModel(
+            d_model=d_model,
+            n_layer=n_layer,
+            vocab_size=vocab_size,
+            initializer_cfg=initializer_cfg,
+            **backbone_kwargs,
+            **factory_kwargs,
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
+
+        # Initialize weights and apply final processing
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+        self.tie_weights()
+        # _tied_weights_keys = ['lm_head.weight']
+
+    def tie_weights(self):
+        self.lm_head.weight = self.backbone.embedding.weight
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+    def forward(
+            self,
+            input_ids,
+            position_ids=None,
+            inference_params=None,
+            num_last_tokens=0,
+            labels=None):
+        """
+        "position_ids" is just to be compatible with Transformer generation. We don't use it.
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+        hidden_states = self.backbone(input_ids, inference_params=inference_params)
+        if num_last_tokens > 0:
+            hidden_states = hidden_states[:, -num_last_tokens:]
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            logits = lm_logits
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+            return (loss,)
+
+        else:
+            CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
+            return CausalLMOutput(logits=lm_logits)