Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -26,118 +26,53 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
     def forward(self, hidden_states: torch.Tensor):
         batch_size, seq_length, hidden_dim = hidden_states.shape
         flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
+        N_tokens = flat_hidden.size(0)
 
+        # 计算路由逻辑
         router_logits = self.gate(flat_hidden)  # (B*seq_len, num_experts)
-        routing_weights = F.softmax(router_logits, dim=1, dtype=flat_hidden.dtype) # (B*seq_len, num_experts)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
 
+        # 选择top-k专家
         routing_weights_topk, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
         if self.norm_topk_prob:
             routing_weights_topk = routing_weights_topk / routing_weights_topk.sum(dim=-1, keepdim=True)
         routing_weights_topk = routing_weights_topk.to(flat_hidden.dtype)
 
-        # ---------- 稀疏计算部分 ----------
-        sparse_output = torch.zeros((flat_hidden.size(0), hidden_dim), 
-                                    dtype=flat_hidden.dtype, device=flat_hidden.device)
-        # 使用 tensor 存储，每个 token 对各专家的输出：形状 (B*seq_len, num_experts, hidden_dim)
-        activated_outputs_tensor = torch.zeros((flat_hidden.size(0), self.num_experts, hidden_dim),
-                                               dtype=flat_hidden.dtype, device=flat_hidden.device)
-        expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts)
-        expert_mask = expert_mask.permute(2, 1, 0)  # (num_experts, top_k, B*seq_len)
-
+        # ---------- 真实计算所有专家输出（密集计算）----------
+        all_expert_outputs = torch.zeros((N_tokens, self.num_experts, hidden_dim),
+                                         dtype=flat_hidden.dtype, device=flat_hidden.device)
+        
         for expert_idx in range(self.num_experts):
             expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx])
-            if top_x.numel() > 0:
-                current_state = flat_hidden[top_x]  # (n, hidden_dim)
-                current_output = expert_layer(current_state)  # (n, hidden_dim)
-                weight = routing_weights_topk[top_x, idx].unsqueeze(-1)  # (n, 1)
-                weighted_output = current_output * weight
-                sparse_output.index_add_(0, top_x, weighted_output.to(flat_hidden.dtype))
-                # 直接存入 tensor：激活 token 对当前专家的输出
-                activated_outputs_tensor[top_x, expert_idx, :] = current_output
-        # ---------- 稀疏计算结束 ----------
-
-        # ---------- Dense估计部分 (向量化版本，激活专家直接使用输出) ----------
-        all_expert_outputs = activated_outputs_tensor  # (B*seq_len, num_experts, hidden_dim)
-        all_routing = selected_experts.tolist()  # list，每个 token 的激活专家列表
+            # 对所有token都计算当前专家的输出
+            expert_output = expert_layer(flat_hidden)  # (N_tokens, hidden_dim)
+            all_expert_outputs[:, expert_idx, :] = expert_output
         
-        N_tokens = flat_hidden.size(0)
-        num_experts = self.num_experts
-
-        # 将 selected_experts 转换为 one-hot 二值矩阵 R: (N_tokens, num_experts)
-        R = F.one_hot(selected_experts, num_classes=num_experts).to(flat_hidden.dtype)  # (N_tokens, top_k, num_experts)
-        R = R.sum(dim=1)  # (N_tokens, num_experts)，激活的专家位置值大于0
-
-        # 计算 token 之间共享激活情况 S: (N_tokens, N_tokens)
-        S = torch.matmul(R, R.t())  # S[i,j] > 0 表示 token i 和 token j 至少共享一个激活专家
-        S = S * (1 - torch.eye(N_tokens, device=S.device, dtype=flat_hidden.dtype))  # 去除自身
-
-        # 构造候选 mask M: (N_tokens, N_tokens, num_experts)
-        # M[i, j, e] = 1 表示 token j 激活了专家 e 且 token i 与 token j 至少共享一个激活专家
-        R_expanded = R.unsqueeze(0).expand(N_tokens, -1, -1)      # (N_tokens, N_tokens, num_experts)
-        S_expanded = S.unsqueeze(-1)                              # (N_tokens, N_tokens, 1)
-        candidate_mask = ((R_expanded > 0) & (S_expanded > 0)).to(flat_hidden.dtype) # (N_tokens, N_tokens, num_experts)
-
-        # 对于数值稳定，排除 token 自身（对角线置0）
-        candidate_mask = candidate_mask * (1 - torch.eye(N_tokens, device=candidate_mask.device, dtype=flat_hidden.dtype).unsqueeze(-1))
-
-        # 扩展 mask 和 all_expert_outputs 以便批量聚合
-        # all_expert_outputs: (N_tokens, num_experts, hidden_dim)
-        candidate_mask_exp = candidate_mask.unsqueeze(-1)          # (N_tokens, N_tokens, num_experts, 1)
-        all_expert_outputs_exp = all_expert_outputs.unsqueeze(0)   # (1, N_tokens, num_experts, hidden_dim)
-
-        # 对每个 token i 和专家 e，聚合候选 token 的输出
-        sum_outputs = (candidate_mask_exp * all_expert_outputs_exp).sum(dim=1)   # (N_tokens, num_experts, hidden_dim)
-        count_outputs = candidate_mask.sum(dim=1).unsqueeze(-1)                  # (N_tokens, num_experts, 1)
-        estimated_dense_all = torch.where(count_outputs > 0, sum_outputs / (count_outputs+1), 
-                                         torch.zeros_like(sum_outputs))  # (N_tokens, num_experts, hidden_dim)
-
-        # 对于激活的专家，直接使用当前 token 的输出
-        # R > 0 表示激活，扩展为 (N_tokens, num_experts, 1) 与 activated_outputs_tensor 对齐
-        activated_mask = (R > 0).unsqueeze(-1)
-        estimated_dense_all = torch.where(activated_mask, activated_outputs_tensor, estimated_dense_all)
-
-        # 利用 gate_prob 加权聚合所有专家输出
-        gate_prob_exp = routing_weights.to(estimated_dense_all.dtype).unsqueeze(-1)  # (N_tokens, num_experts, 1)
-        dense_outputs = (gate_prob_exp * estimated_dense_all).sum(dim=1)  # (N_tokens, hidden_dim)
-        # ---------- Dense估计结束 (向量化版本) ----------
-
+        # ---------- 提取激活专家输出（稀疏前向）----------
+        # 计算稀疏输出
+        sparse_output = torch.zeros((N_tokens, hidden_dim), 
+                                    dtype=flat_hidden.dtype, device=flat_hidden.device)
+        
+        # 为每个token，提取并加权其激活专家的输出
+        for token_idx in range(N_tokens):
+            for k in range(self.top_k):
+                expert_idx = selected_experts[token_idx, k].item()
+                weight = routing_weights_topk[token_idx, k]
+                sparse_output[token_idx] += all_expert_outputs[token_idx, expert_idx] * weight
+        
+        # ---------- 密集计算聚合（用于反向传播）----------
+        # 使用所有专家的输出和路由权重计算密集输出
+        routing_weights_expanded = routing_weights.unsqueeze(-1)  # (N_tokens, num_experts, 1)
+        dense_outputs = (all_expert_outputs * routing_weights_expanded).sum(dim=1)  # (N_tokens, hidden_dim)
+        
+        # ---------- 组合稀疏前向和密集反向 ----------
+        # sparse_output.detach()保留稀疏前向计算图
+        # (dense_outputs - dense_outputs.detach())只保留密集反向梯度
         final_flat = sparse_output.detach() + (dense_outputs - dense_outputs.detach())
         final_output = final_flat.view(batch_size, seq_length, hidden_dim)
+        
         return final_output, router_logits
 
-    def estimate_dense_output(self, token_idx, activated, gate_prob, activated_outputs, all_routing, all_expert_outputs):
-        num_experts = gate_prob.size(0)
-        dense_parts = {}
-        # 对于激活的专家，直接使用 tensor 的对应行
-        for idx in activated:
-            dense_parts[idx] = activated_outputs[idx]
-        non_activated = [i for i in range(num_experts) if i not in activated]
-        for i in non_activated:
-            indices = []
-            for idx, r_dec in enumerate(all_routing):
-                if (i in r_dec) and (len(set(r_dec) & set(activated)) > 0):
-                    indices.append(idx)
-            if indices:
-                selected_outputs = all_expert_outputs[indices, i, :]  # (n, hidden_dim)
-                mask = (selected_outputs.sum(dim=-1) != 0).to(selected_outputs.dtype).unsqueeze(-1)
-                if mask.sum() > 0:
-                    estimated = (selected_outputs * mask).sum(dim=0) / mask.sum()
-                else:
-                    estimated = torch.zeros_like(selected_outputs[0])
-            else:
-                all_outputs = all_expert_outputs[:, i, :]
-                mask = (all_outputs.sum(dim=-1) != 0).to(all_outputs.dtype).unsqueeze(-1)
-                if mask.sum() > 0:
-                    estimated = (all_outputs * mask).sum(dim=0) / mask.sum()
-                else:
-                    estimated = torch.zeros_like(all_outputs[0])
-            dense_parts[i] = estimated
-        estimated_dense = 0
-        for i in range(num_experts):
-            estimated_dense += gate_prob[i] * dense_parts[i]
-        return estimated_dense
-
 class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     """
     自定义的 Olmoe ForCausalLM 模型，使用新的 DenseBackwardOlmoeSparseMoeBlock 替换原版的 MoE 模块，