Update config.json

config.json

@@ -1,52 +1,25 @@
 {
-   "name":"LayeredMambaLM",
-   "input":{
-      "vocab_size":128256,
-      "tie_embeddings":true,
-      "pad_vocab_size_multiple":8,
-      "lm_head_bias":false
-   },
-   "MixerModel":{
-      "name":"MixerModel",
-      "input":{
-         "d_model":2048,
-         "n_layer":16,
-         "lm_head_prenorm":"rms"
-      },
-      "Blocks":[
-         {
-            "name":"LlamaBlock",
-            "n_layers":16,
-            "input":{
-               "resid_dropout":0.0,
-               "mlp_intermediate_size":8192,
-               "mlp_act_fn":"silu"
-            },
-            "Layer":{
-               "name":"DiscreteMamba2",
-               "input":{
-                  "d_state":64,
-                  "n_qk_heads":32,
-                  "n_v_heads":32,
-                  "expand":1,
-                  "chunk_size":128,
-                  "activation":"identity",
-                  "use_ref_impl":false,
-                  "bias":false,
-                  "norm_cls":"none",
-                  "initializer":{
-                     "a_log":"default",
-                     "x":"default",
-                     "B":"default",
-                     "C":"default",
-                     "D":"default",
-                     "z":"identity",
-                     "out":"default",
-                     "convolution":"identity"
-                  }
-               }
-            }
-         }
-      ]
-   }
-}
+  "model_type": "llamba",
+  "vocab_size": 128256,
+  "tie_embeddings": true,
+  "pad_vocab_size_multiple": 8,
+  "lm_head_bias": false,
+  "d_model": 2048,
+  "n_layer": 32,
+  "resid_dropout": 0.0,
+  "norm_epsilon": 1e-5,
+  "mlp_cfg": {
+    "intermediate_size": 14336,
+    "bias": false,
+    "act_fn": "silu"
+  },
+  "ssm_cfg": {
+    "d_state": 64,
+    "n_v_heads": 32,
+    "n_qk_heads": 32,
+    "expand": 1,
+    "chunk_size": 128,
+    "activation": "identity",
+    "bias": false
+  }
+}