config shenanigans

myolmoe/config.json

@@ -15,8 +15,8 @@
   "norm_topk_prob": false,
   "num_attention_heads": 16,
   "num_experts": 64,
-  "num_small_experts": 64,
-  "small_expert_intermediate_size": 512,
+  "num_small_experts": 0,
+  "small_expert_intermediate_size": 0,
   "num_experts_per_tok": 2,
   "num_hidden_layers": 16,
   "num_key_value_heads": 16,

scripts/train.py

@@ -41,6 +41,8 @@ def expand_model_with_small_experts(base_model):
     print("# DEBUG: Expanding model with small experts...")
     config = base_model.config
     expanded_model = MyOlmoeForCausalLM(config)
+    config.num_small_experts = 64  # Add 64 small experts
+    config.small_expert_intermediate_size = config.intermediate_size // 16  # Half size
     
     base_state_dict = base_model.state_dict()
     expanded_state_dict = expanded_model.state_dict()