Update configuration.py

configuration.py

@@ -1,25 +1,30 @@
 from transformers import PretrainedConfig
 
-class MoLMConfig(PretrainedConfig):
-    model_type = "MoLM"
+class MoEGPTConfig(PretrainedConfig):
+    model_type = "moegpt"
 
     def __init__(
         self,
         vocab_size=50304,
-        n_embd=768,
-        n_layer=12,
-        n_head=12,
+        n_embd=1152, #768,
+        n_layer=24,#12,
+        n_head=16,#12,
         sequence_length=1024,
+        moe=False,
+        moe_routing="standard_gating",
+        moe_num_experts=4,
+        moe_num_experts_per_tok=2,
+        moe_softmax_order="softmax_topk",
+        moe_router_loss="load_balancing_z_loss",
+        moe_aux_loss_factor=0.01,
+        moe_z_loss_factor=1.0,
         mlp_dim_exp_factor=1.0,
         dropout=0.0,
         bias=False,
-        num_experts=6,
-        expert_configs=None,
-        use_router=False,
-        architectures=["MoLM"],
+        architectures=["MoEGPTForCausalLM"],
         auto_map={
-            "AutoConfig": "configuration.MoLMConfig",
-            "AutoModelForCausalLM": "modeling.MoLM",
+            "AutoConfig": "configuration.MoEGPTConfig",
+            "AutoModelForCausalLM": "modeling.MoEGPTForCausalLM",
             "AutoTokenizer": "GPT2TokenizerFast"
         },
         **kwargs,
@@ -30,31 +35,17 @@ class MoLMConfig(PretrainedConfig):
         self.n_layer = n_layer
         self.n_head = n_head
         self.sequence_length = sequence_length
+        self.moe = moe
+        self.moe_routing = moe_routing
+        self.moe_num_experts = moe_num_experts
+        self.moe_num_experts_per_tok = moe_num_experts_per_tok
+        self.moe_softmax_order = moe_softmax_order
+        self.moe_router_loss = moe_router_loss
+        self.moe_aux_loss_factor = moe_aux_loss_factor
+        self.moe_z_loss_factor = moe_z_loss_factor
         self.mlp_dim_exp_factor = mlp_dim_exp_factor
         self.dropout = dropout
         self.bias = bias
-        self.num_experts = num_experts
-        self.expert_configs = expert_configs
-        self.use_router = use_router
         self.architectures = architectures
         self.auto_map = auto_map
-
-    def to_dict(self):
-        config_dict = super().to_dict()
-        config_dict.update({
-            "vocab_size": self.vocab_size,
-            "n_embd": self.n_embd,
-            "n_layer": self.n_layer,
-            "n_head": self.n_head,
-            "sequence_length": self.sequence_length,
-            "mlp_dim_exp_factor": self.mlp_dim_exp_factor,
-            "dropout": self.dropout,
-            "bias": self.bias,
-            "num_experts": self.num_experts,
-            "expert_configs": [
-                expert_config.to_dict() if not isinstance(expert_config, dict) else expert_config 
-                for expert_config in self.expert_configs
-            ] if self.expert_configs else None,  
-            "use_router": self.use_router
-        })
-        return config_dict
+