Update config.json

config.json

@@ -1 +1,25 @@
-{"name": "LayeredMambaLM", "input": {"vocab_size": 128256, "tie_embeddings": false, "pad_vocab_size_multiple": 8, "lm_head_bias": false}, "MixerModel": {"name": "MixerModel", "input": {"d_model": 4096, "n_layer": 32, "lm_head_prenorm": "rms"}, "Blocks": [{"name": "LlamaBlock", "n_layers": 32, "input": {"resid_dropout": 0.0, "mlp_intermediate_size": 14336, "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": false,
+  "pad_vocab_size_multiple": 8,
+  "lm_head_bias": false,
+  "d_model": 4096,
+  "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
+  }
+}