Upload DogeForCausalLM

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "/root/autodl-tmp/data/Doge-320M",
+  "_name_or_path": "/root/autodl-tmp/small-doge/data/Doge-320M-decay/checkpoint-4000",
   "architectures": [
     "DogeForCausalLM"
   ],
@@ -11,7 +11,6 @@
   "bos_token_id": 0,
   "dynamic_mask_ratio": 0.0,
   "eos_token_id": 1,
-  "expert_retrieval_size": 64,
   "hidden_act": "silu",
   "hidden_bias": false,
   "hidden_dropout": 0.0,

configuration_doge.py

@@ -121,8 +121,6 @@ class DogeConfig(PretrainedConfig):
             Number of Experts for the Cross Domain Mixture of Experts.
         num_experts_per_tok (`int`, *optional*, defaults to 8):
             Number of selected experts to route per-token.
-        expert_retrieval_size (`int`, *optional*, defaults to 64):
-            Dimension of the Expert retrieval states for calculating the dot product of query and key to determine the expert index.
 
     ```python
     >>> from transformers import DogeConfig, DogeModel
@@ -149,7 +147,7 @@ class DogeConfig(PretrainedConfig):
         "layers.*.feed_forward.gate_proj": "colwise",
         "layers.*.feed_forward.up_proj": "colwise",
         "layers.*.feed_forward.down_proj": "rowwise",
-        "layers.*.feed_forward.queries_proj": "colwise",
+        "layers.*.feed_forward.router_gate": "colwise",
         "layers.*.feed_forward.down_embed": "rowwise",
         "layers.*.feed_forward.up_embed": "rowwise",
     }
@@ -181,7 +179,6 @@ class DogeConfig(PretrainedConfig):
         is_moe=False,
         num_experts=2048,
         num_experts_per_tok=8,
-        expert_retrieval_size=64,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -207,7 +204,6 @@ class DogeConfig(PretrainedConfig):
         self.is_moe = is_moe
         self.num_experts = num_experts
         self.num_experts_per_tok = num_experts_per_tok
-        self.expert_retrieval_size = expert_retrieval_size
 
         # Validate the correctness of rotary position embeddings parameters
         # BC: if there is a 'type' field, copy it it to 'rope_type'.

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bb05e3cf42d7df4c683d7c6719d195d06614686b766ed0782bc2f3b7c71afec5
+oid sha256:ce4aaf436761b12719bb9be9d3a250ba388679b324886299ed71f69c2b53a510
 size 1343277696

modeling_doge.py

@@ -480,23 +480,17 @@ class DogeCDMoE(DogeMLP):
         self.hidden_dim = config.hidden_size
         self.act_fn = ACT2FN[config.hidden_act]
 
-        self.expert_retrieval_dim = config.expert_retrieval_size
         self.num_experts = config.num_experts
         self.top_k = config.num_experts_per_tok
         self.num_keys = int(math.sqrt(self.num_experts))
 
-        # queries and keys for retrieval experts
-        self.queries_proj = nn.Linear(self.hidden_dim, self.expert_retrieval_dim, bias=False)
-        self.keys = nn.Parameter(torch.zeros(2, self.expert_retrieval_dim // 2, self.num_keys))
+        # router gate for retrieval experts
+        self.router_gate = nn.Linear(self.hidden_dim, self.num_keys * 2)
 
         # experts
         self.down_embed = nn.Embedding(self.num_experts, self.hidden_dim)
         self.up_embed = nn.Embedding(self.num_experts, self.hidden_dim)
 
-        # scaling factor
-        self.mlp_scaling = nn.Parameter(torch.ones(self.hidden_dim))
-        self.moe_scaling = nn.Parameter(torch.zeros(self.hidden_dim))
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -504,27 +498,25 @@ class DogeCDMoE(DogeMLP):
     ) -> torch.Tensor:
         bsz, seq_len, _ = hidden_states.shape
 
-        # get routing weights with queries and keys
-        queries = self.queries_proj(hidden_states).view(2, bsz * seq_len, -1)
-        routing_weights = torch.matmul(queries, self.keys)
+        # get routing weights with router gate
+        routing_weights = self.router_gate(hidden_states).view(2, bsz * seq_len, -1)
 
         # get experts with the highest routing weights
-        (scores_x, scores_y), (indices_x, indices_y) = routing_weights.topk(self.top_k, dim=-1)
+        (scores_x, scores_y), (indices_x, indices_y) = [w.topk(self.num_keys, dim=-1) for w in routing_weights]
         all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)
-        all_scores = all_scores.view(*scores_x.shape[:-1], -1)
-        all_indices = (indices_x.unsqueeze(-1) * self.num_keys) + indices_y.unsqueeze(-2)
-        all_indices = all_indices.view(*indices_x.shape[:-1], -1)
-        scores, pk_indices = all_scores.topk(self.top_k, dim=-1)
-        indices = all_indices.gather(-1, pk_indices)
-        down_embed = self.down_embed(indices).transpose(1, 2)
+        all_indices = indices_x.unsqueeze(-1) * self.num_keys + indices_y.unsqueeze(-2)
+        all_scores = all_scores.view(*all_scores.shape[:-2], -1)
+        all_indices = all_indices.view(*all_indices.shape[:-2], -1)
+        scores, indices = all_scores.topk(self.top_k, dim=-1)
+        down_embed = self.down_embed(indices)
         up_embed = self.up_embed(indices)
 
         # mix experts states with cross domain states
-        experts_weights = torch.matmul(hidden_states.view(bsz * seq_len, 1, -1), down_embed).view(bsz * seq_len, -1)
+        experts_weights = torch.matmul(down_embed, hidden_states.view(bsz * seq_len, -1, 1)).view(bsz * seq_len, -1)
         experts_weights = self.act_fn(experts_weights) * scores.softmax(dim=-1)
         experts_states = torch.matmul(experts_weights.view(bsz * seq_len, 1, -1), up_embed).view(bsz, seq_len, -1)
         hidden_states = self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
-        hidden_states = (hidden_states * self.mlp_scaling) + (experts_states * self.moe_scaling)
+        hidden_states = hidden_states + experts_states
         return hidden_states