Update config.json

config.json

@@ -1 +1,52 @@
-{"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"}}}}]}}
+{
+   "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"
+                  }
+               }
+            }
+         }
+      ]
+   }
+}