Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -23,9 +23,9 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         final_output: Tensor, shape (batch_size, sequence_length, hidden_dim)
         router_logits: Tensor, shape (batch_size * sequence_length, num_experts)
     """
-    def forward(self, hidden_states: torch.Tensor):
+    def forward_partscale_fixep_norm_dtch(self, hidden_states: torch.Tensor):
         """
-        forward_partscale_fixep_norm_dtch_tau
+        forward_partscale_fixep_norm_dtch
         """
         batch_size, seq_length, hidden_dim = hidden_states.shape
         dtype = hidden_states.dtype
@@ -37,22 +37,17 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         # Compute routing logic
         router_logits = self.gate(flat_hidden).to(dtype=dtype)  # (B*L, num_experts)
         routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*L, num_experts)
-        routing_weights_tau = F.softmax(router_logits / 1.1, dim=1, dtype=torch.float)  # (B*L, num_experts)
 
         # Select top-k experts
         routing_weights_topk, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
-        routing_weights_topk_tau, selected_experts_tau = torch.topk(routing_weights_tau, self.top_k, dim=-1)
         if self.norm_topk_prob:
             norm_ratio = routing_weights_topk.sum(dim=-1, keepdim=True)
             # Normalize top-k routing weights
             routing_weights_topk = routing_weights_topk / norm_ratio
             # Only scale the selected top-k positions in routing_weights
-            mask = F.one_hot(selected_experts_tau, num_classes=self.num_experts).sum(dim=1).to(dtype)
-            routing_weights_topk_tau = routing_weights_tau * mask
-            norm_ratio_dense = routing_weights_topk_tau.sum(dim=-1, keepdim=True)
+            mask = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1).to(dtype)
             # ------------------------------------Choose Section-----------------------------------------------
             # current --> partscale_fix_expert implementation
-            routing_weights_tau = routing_weights_tau * (1.0 - mask) / norm_ratio_dense.detach() + routing_weights_topk_tau / norm_ratio_dense
             routing_weights = routing_weights * (1.0 - mask) / norm_ratio.detach() + routing_weights * mask / norm_ratio
 
             # should be --> the gated implemenation, by comment out the line above and uncomment the two lines below
@@ -63,8 +58,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         routing_weights_topk = routing_weights_topk.to(dtype=dtype)
 
         # Convert full routing_weights to consistent dtype for dense accumulation
-        # routing_weights = routing_weights.to(dtype=dtype)
-        routing_weights_tau = routing_weights_tau.to(dtype=dtype)
+        routing_weights = routing_weights.to(dtype=dtype)
 
         # Prepare accumulators: one for dense_outputs, one for sparse_outputs
         dense_outputs = torch.zeros((N_tokens, hidden_dim), dtype=dtype, device=device)
@@ -77,14 +71,13 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
             expert_layer = self.experts[expert_idx]
             # Compute current expert output for all tokens
             expert_output = expert_layer(flat_hidden).to(dtype=dtype)  # (N_tokens, hidden_dim)
-            activation_mask = (selected_experts_tau == expert_idx).any(dim=1).float().unsqueeze(-1).to(dtype)
+            activation_mask = (selected_experts == expert_idx).any(dim=1).float().unsqueeze(-1).to(dtype)
             if expert_output.requires_grad:
                 expert_output.register_hook(lambda grad, mask=activation_mask: grad * mask)
             expert_output = expert_output.to(dtype=dtype)
             # Dense accumulation: multiply by full routing weight and add
-            weight_full_tau = routing_weights_tau[:, expert_idx].unsqueeze(-1)  # (N_tokens, 1)
             weight_full = routing_weights[:, expert_idx].unsqueeze(-1)  # (N_tokens, 1)
-            dense_outputs = dense_outputs + expert_output * (weight_full_tau-weight_full_tau.detach()) + expert_output * weight_full.detach()
+            dense_outputs = dense_outputs + expert_output * weight_full
 
             # Sparse accumulation: find tokens where this expert is among top_k
             # matches: Boolean mask where selected_experts == expert_idx → shape (N_tokens, top_k)