Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -25,7 +25,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
     """
     def forward(self, hidden_states: torch.Tensor):
         """
-        Gate version of implementation of straight-through, π -> mask, dmask / dπ = 1
+        forward_partscale_fixep_norm_dtch_tau
         """
         batch_size, seq_length, hidden_dim = hidden_states.shape
         dtype = hidden_states.dtype
@@ -34,49 +34,73 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
         N_tokens = flat_hidden.size(0)
 
-        # 1) router & softmax
-        router_logits = self.gate(flat_hidden).to(dtype=dtype)                # (N, num_experts)
-        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (N, num_experts)
+        # 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)
 
-        # 2) top-K selection
-        _, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)  # (N, K), (N, K)
-        
-        # 3) build hard & ste masks
-        mask_hard = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1).to(dtype)     # (N, num_experts)
-        #mask_ste = mask_hard + (routing_weights - routing_weights.detach())
-        mask_ste = mask_hard + (router_logits - router_logits.detach())
-
-        # 4) compute gated weights = π * mask, then optionally renormalize
-        gated = routing_weights * mask_ste                              # zero-out non-TopK
+        # 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 = gated.sum(dim=-1, keepdim=True)                # (N,1)
-            gated = gated / norm_ratio                                  # normalized TopK
+            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)
+            # ------------------------------------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
+            # gated = routing_weights.detach() * mask + (routing_weights - routing_weights.detach())
+            # routing_weights = gated / gated.sum(dim=-1, keepdim=True)
+            # ------------------------------------Choose Section-----------------------------------------------
 
-        # 5）prepare accumulators
+        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)
+
+        # Prepare accumulators: one for dense_outputs, one for sparse_outputs
         dense_outputs = torch.zeros((N_tokens, hidden_dim), dtype=dtype, device=device)
         sparse_outputs = torch.zeros((N_tokens, hidden_dim), dtype=dtype, device=device)
 
-        for expert_idx, expert_layer in enumerate(self.experts):
-            expert_output = expert_layer(flat_hidden).to(dtype=dtype)  # (N_tokens, hidden_dim)
-            activation_mask = (selected_experts == expert_idx).any(dim=1).float().unsqueeze(-1).to(dtype)
+        # For mapping top-k positions when accumulating sparse_outputs
+        # selected_experts: (N_tokens, top_k)
 
+        for expert_idx in range(self.num_experts):
+            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)
             if expert_output.requires_grad:
                 expert_output.register_hook(lambda grad, mask=activation_mask: grad * mask)
-            
-            # a) Dense-STE backward uses gated weights
-            weights = gated[:, expert_idx].unsqueeze(-1)  # (N_tokens, 1)
-            dense_outputs += expert_output * weights
-
-            # b) Sparse forward -- find tokens where this expert is among top_k (active experts)
-            active = (selected_experts == expert_idx)
-            if active.any():
-                token_indices, _ = torch.where(active)
-                weights_topk = gated[token_indices, expert_idx].unsqueeze(-1)  # (num_matches,1)
-                sparse_outputs[token_indices] += expert_output[token_indices] * weights_topk
-        
-        # 6) STE mix: forward from sparse, backward from dense
+            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()
+
+            # Sparse accumulation: find tokens where this expert is among top_k
+            # matches: Boolean mask where selected_experts == expert_idx → shape (N_tokens, top_k)
+            matches = (selected_experts == expert_idx)
+            if matches.any():
+                # locations: tuple of (token_indices, k_indices)
+                token_indices, k_indices = torch.where(matches)
+                # corresponding top-k weights
+                weights_topk = routing_weights_topk[token_indices, k_indices].unsqueeze(-1)  # (num_matches, 1)
+                # Accumulate sparse_outputs only for matched tokens
+                sparse_outputs[token_indices] = sparse_outputs[token_indices] + expert_output[token_indices] * weights_topk
+
+        # Combine sparse forward output and dense backward output
         final_flat = sparse_outputs.detach() + (dense_outputs - dense_outputs.detach())
-        final_output = final_flat.view(batch_size, seq_length, hidden_dim).to(dtype=dtype)
+        final_flat = final_flat.to(dtype=dtype)
+        final_output = final_flat.view(batch_size, seq_length, hidden_dim)
 
         return final_output, router_logits
 
@@ -85,6 +109,71 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
 
 
 
+
+
+    
+    # def forward(self, hidden_states: torch.Tensor):
+    #     """
+    #     Gate version of implementation of straight-through, π -> mask, dmask / dπ = 1
+    #     """
+    #     batch_size, seq_length, hidden_dim = hidden_states.shape
+    #     dtype = hidden_states.dtype
+    #     device = hidden_states.device
+
+    #     flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
+    #     N_tokens = flat_hidden.size(0)
+
+    #     # 1) router & softmax
+    #     router_logits = self.gate(flat_hidden).to(dtype=dtype)                # (N, num_experts)
+    #     routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (N, num_experts)
+
+    #     # 2) top-K selection
+    #     _, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)  # (N, K), (N, K)
+        
+    #     # 3) build hard & ste masks
+    #     mask_hard = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1).to(dtype)     # (N, num_experts)
+    #     #mask_ste = mask_hard + (routing_weights - routing_weights.detach())
+    #     mask_ste = mask_hard + (router_logits - router_logits.detach())
+
+    #     # 4) compute gated weights = π * mask, then optionally renormalize
+    #     gated = routing_weights * mask_ste                              # zero-out non-TopK
+    #     if self.norm_topk_prob:
+    #         norm_ratio = gated.sum(dim=-1, keepdim=True)                # (N,1)
+    #         gated = gated / norm_ratio                                  # normalized TopK
+
+    #     # 5）prepare accumulators
+    #     dense_outputs = torch.zeros((N_tokens, hidden_dim), dtype=dtype, device=device)
+    #     sparse_outputs = torch.zeros((N_tokens, hidden_dim), dtype=dtype, device=device)
+
+    #     for expert_idx, expert_layer in enumerate(self.experts):
+    #         expert_output = expert_layer(flat_hidden).to(dtype=dtype)  # (N_tokens, hidden_dim)
+    #         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)
+            
+    #         # a) Dense-STE backward uses gated weights
+    #         weights = gated[:, expert_idx].unsqueeze(-1)  # (N_tokens, 1)
+    #         dense_outputs += expert_output * weights
+
+    #         # b) Sparse forward -- find tokens where this expert is among top_k (active experts)
+    #         active = (selected_experts == expert_idx)
+    #         if active.any():
+    #             token_indices, _ = torch.where(active)
+    #             weights_topk = gated[token_indices, expert_idx].unsqueeze(-1)  # (num_matches,1)
+    #             sparse_outputs[token_indices] += expert_output[token_indices] * weights_topk
+        
+    #     # 6) STE mix: forward from sparse, backward from dense
+    #     final_flat = sparse_outputs.detach() + (dense_outputs - dense_outputs.detach())
+    #     final_output = final_flat.view(batch_size, seq_length, hidden_dim).to(dtype=dtype)
+
+    #     return final_output, router_logits
+
+
+
+
+
+
     # def forward(self, hidden_states: torch.Tensor):
     #     batch_size, seq_length, hidden_dim = hidden_states.shape
     #     # 记录输入张量的数据类型，确保所有计算保持一致