Update modeling_minicpm.py

modeling_minicpm.py

@@ -304,49 +304,140 @@ class AddAuxiliaryLoss(torch.autograd.Function):
 
 
 class MiniCPMMoE(nn.Module):
-    def __init__(self, config):
+    """
+    MiniCPM MoE with Default MoE implementation.
+
+    Based on paper: "Dense Backpropagation Improves Training for Sparse Mixture-of-Experts"
+
+    Key idea:
+    - Sparse forward: only compute top-K experts
+    - Dense backward: router receives gradients from ALL experts via default vectors
+    - EMA update: default vectors are updated with exponential moving average of expert outputs
+    """
+    def __init__(self, config, beta=0.9):
         super().__init__()
         self.config = config
         self.num_experts = config.num_experts
         self.num_experts_per_tok = config.num_experts_per_tok
+        self.hidden_size = config.hidden_size
         self.experts = nn.ModuleList(
             [MiniCPMMLP(config) for i in range(self.num_experts)]
         )
         self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
         self.intermediate_size = config.intermediate_size
-    
+
+        # Default MoE: EMA parameter and default vectors
+        self.beta = beta  # EMA decay coefficient
+
+        # Register default vector buffer for each expert
+        for expert_idx in range(self.num_experts):
+            self.register_buffer(
+                f'default_vector_{expert_idx}',
+                torch.zeros(config.hidden_size)
+            )
+
     def forward(self, hidden_states):
+        """
+        Default MoE forward pass.
+
+        Algorithm (from paper "Dense Backpropagation Improves Training for Sparse Mixture-of-Experts"):
+        1. Compute routing weights: π = Softmax(W·x)
+        2. Select top-K experts: A = TopK(π)
+        3. Compute output:
+           y = Σ πi · { Ei(x)  if i ∈ A
+                      { Êi     if i ∉ A
+        4. Update EMA: Êi^(t) = β·Êi^(t-1) + (1-β)·mean(Ei(x)) for activated experts
+
+        Key advantages:
+        - Sparse forward: only top-K experts are computed
+        - Dense backward: router receives gradient signals from all N experts
+        """
         orig_shape = hidden_states.shape
         orig_dtype = hidden_states.dtype
-        hidden_states = hidden_states.view(-1, orig_shape[-1])
-        token_num = hidden_states.shape[0]
-        scores = self.gate(hidden_states)
-        scores_prob = F.softmax(scores, dim=-1, dtype=torch.float32)
-        expert_weights, expert_indices = torch.topk(scores_prob, self.num_experts_per_tok, dim=-1)
-        expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True)
-        topk_idx_flat = expert_indices.view(-1)
-        expert_weights = expert_weights.to(orig_dtype)
-
+        device = hidden_states.device
+
+        flat_hidden = hidden_states.view(-1, orig_shape[-1])  # (N_tokens, hidden_dim)
+        N_tokens = flat_hidden.shape[0]
+        hidden_dim = orig_shape[-1]
+
+        # ========== Step 1: Compute routing weights ==========
+        scores = self.gate(flat_hidden)  # (N_tokens, num_experts)
+        scores_prob = F.softmax(scores, dim=-1, dtype=torch.float32)  # (N_tokens, num_experts)
+
+        # ========== Step 2: Select top-K experts ==========
+        expert_weights_topk, expert_indices = torch.topk(scores_prob, self.num_experts_per_tok, dim=-1)
+        # (N_tokens, top_k), (N_tokens, top_k)
+
+        # Top-K normalization
+        expert_weights_topk = expert_weights_topk / expert_weights_topk.sum(dim=-1, keepdim=True)
+        expert_weights_topk = expert_weights_topk.to(orig_dtype)
+        scores_prob = scores_prob.to(orig_dtype)
+
+        # ========== Step 3: Compute expert outputs (sparse + default vectors) ==========
+        final_output = torch.zeros((N_tokens, hidden_dim), dtype=orig_dtype, device=device)
+
+        for expert_idx in range(self.num_experts):
+            expert_layer = self.experts[expert_idx]
+
+            # Get default vector for this expert
+            default_vector = getattr(self, f'default_vector_{expert_idx}').to(dtype=orig_dtype)
+
+            # Find which tokens activated this expert
+            matches = (expert_indices == expert_idx)  # (N_tokens, top_k)
+            is_activated = matches.any(dim=1)  # (N_tokens,)
+
+            if is_activated.any():
+                # ===== Activated tokens: compute real output =====
+                activated_token_indices = torch.where(is_activated)[0]
+                activated_inputs = flat_hidden[activated_token_indices]  # (n_activated, hidden_dim)
+
+                # Compute real expert output
+                real_expert_output = expert_layer(activated_inputs)  # (n_activated, hidden_dim)
+                real_expert_output = real_expert_output.to(dtype=orig_dtype)
+
+                # ===== Update EMA for this expert (only during training) =====
+                if self.training:
+                    # Compute mean output of activated tokens
+                    mean_output = real_expert_output.mean(dim=0).detach()  # (hidden_dim,)
+                    # EMA update: Êi^(t) = β·Êi^(t-1) + (1-β)·mean(Ei(x))
+                    new_default = self.beta * default_vector + (1 - self.beta) * mean_output
+                    getattr(self, f'default_vector_{expert_idx}').copy_(new_default)
+
+                # ===== Accumulate real output for activated tokens (using normalized top-K weights) =====
+                token_indices, k_indices = torch.where(matches)
+                if len(token_indices) > 0:
+                    # Get corresponding weights
+                    weights = expert_weights_topk[token_indices, k_indices, None]  # (n_matches, 1)
+                    weighted_output = real_expert_output * weights  # (n_matches, hidden_dim)
+
+                    # Efficient accumulation using index_add_
+                    final_output.index_add_(0, token_indices, weighted_output.to(orig_dtype))
+
+            # ===== Non-activated tokens: accumulate default vector (using original softmax weights) =====
+            non_activated_indices = torch.where(~is_activated)[0]
+            if len(non_activated_indices) > 0:
+                # Get routing weights for non-activated tokens (original softmax, not normalized)
+                weights_non_activated = scores_prob[non_activated_indices, expert_idx].unsqueeze(-1)  # (n_non, 1)
+                # Accumulate: weight * default_vector
+                final_output[non_activated_indices] += weights_non_activated * default_vector
+
+        # ========== Step 4: Compute load balancing loss (only during training) ==========
         if self.training:
-            hidden_states = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0)
-            y = torch.empty_like(hidden_states)
-            for i in range(self.num_experts):
-                y[topk_idx_flat == i] = self.experts[i](hidden_states[topk_idx_flat == i])
-            y = (y.view(*expert_weights.shape, -1) * expert_weights.unsqueeze(-1)).sum(dim=1)
-            y =  y.view(*orig_shape)
-
             load = expert_indices.view(-1).bincount(minlength=self.num_experts)
-            load_mean = load / (token_num * self.num_experts_per_tok)
+            load_mean = load / (N_tokens * self.num_experts_per_tok)
             importance_mean = scores_prob.mean(dim=0)
             balance_loss = self.num_experts * torch.sum(importance_mean * load_mean)
 
-            y = AddAuxiliaryLoss.apply(y, balance_loss)
-        else:
-            y = self.moe_infer(hidden_states, topk_idx_flat, expert_weights.view(-1, 1)).view(*orig_shape)
-        return y
-    
+            final_output = AddAuxiliaryLoss.apply(final_output, balance_loss)
+
+        # ========== Step 5: Reshape back to original shape ==========
+        final_output = final_output.view(*orig_shape)
+
+        return final_output
+
     @torch.no_grad()
     def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
+        """Original inference method (not used in Default MoE, kept for compatibility)"""
         expert_cache = torch.zeros_like(x)
         idxs = flat_expert_indices.argsort()
         tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)