Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -14,7 +14,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         batch_size, seq_length, hidden_dim = hidden_states.shape
         flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
 
-        # 计算路由 logits 和 routing 权重
+        # 计算路由 logits 和路由权重
         router_logits = self.gate(flat_hidden)  # (B*seq_len, num_experts)
         routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
 
@@ -28,120 +28,127 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         # 初始化稀疏输出
         sparse_output = torch.zeros((flat_hidden.size(0), hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
         
-        # 存储所有激活信息的数据结构
-        num_tokens = flat_hidden.size(0)
-        all_activated_outputs = {}  # {expert_idx: {token_idx: output_tensor}}
-        all_routing_indices = {}    # {expert_idx: [token_indices]}
-        token_activated_experts = {}  # {token_idx: [activated_expert_indices]}
+        # 记录每个专家被激活的token
+        token_indices_per_expert = [[] for _ in range(self.num_experts)]
+        expert_outputs_dict = {i: {} for i in range(self.num_experts)}
         
         # one-hot 编码 top-k 专家
         expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts)  # (B*seq_len, top_k, num_experts)
         expert_mask = expert_mask.permute(2, 1, 0)  # (num_experts, top_k, B*seq_len)
 
-        # 稀疏计算，同时记录激活情况
+        # 保存每个token激活的专家列表
+        token_to_experts = [[] for _ in range(flat_hidden.size(0))]
+        
+        # 稀疏计算
         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:
+                # 记录该专家被哪些token激活
+                top_x_list = top_x.tolist()
+                token_indices_per_expert[expert_idx].extend(top_x_list)
+                
+                # 记录每个token激活了哪些专家
+                for token_idx in top_x_list:
+                    token_to_experts[token_idx].append(expert_idx)
+                
+                # 标准MoE前向计算
                 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))
                 
-                # 记录该专家激活的token和对应输出
-                all_activated_outputs[expert_idx] = {}
-                all_routing_indices[expert_idx] = top_x.tolist()
-                
-                for pos, token_idx in enumerate(top_x.tolist()):
-                    # 记录该专家对该token的输出
-                    all_activated_outputs[expert_idx][token_idx] = current_output[pos]
-                    
-                    # 记录该token激活的专家
-                    if token_idx not in token_activated_experts:
-                        token_activated_experts[token_idx] = []
-                    token_activated_experts[token_idx].append(expert_idx)
-        # ---------- 稀疏计算结束 ----------
-
-        # ---------- Dense估计部分 ----------
-        # 将activated_experts 转换为list格式，与路由权重匹配
-        all_routing = selected_experts.tolist()  # 长度为 (B*seq_len)
-        
-        # 使用已激活信息估计dense输出
-        dense_outputs = []
-        for token_idx in range(num_tokens):
-            # 获取当前token的激活专家列表
-            activated = all_routing[token_idx] if token_idx in token_activated_experts else []
-            
-            # 估计dense输出（只使用已经计算过的专家输出）
-            dense_est = self.estimate_dense_output_efficient(
-                token_idx=token_idx,
-                activated=activated,
-                gate_prob=routing_weights[token_idx],
-                all_activated_outputs=all_activated_outputs,
-                all_routing_indices=all_routing_indices,
-                token_activated_experts=token_activated_experts
-            )
-            dense_outputs.append(dense_est.unsqueeze(0))
+                # 保存未加权的专家输出，用于后续估计
+                for i, token_idx in enumerate(top_x_list):
+                    expert_outputs_dict[expert_idx][token_idx] = current_output[i]
+        
+        # ---------- 初始化密集估计输出 ----------
+        # 创建一个每个token对所有专家输出的张量
+        dense_outputs = torch.zeros((flat_hidden.size(0), self.num_experts, hidden_dim), 
+                                   dtype=flat_hidden.dtype, device=flat_hidden.device)
+        
+        # 首先将���计算的专家输出填入
+        for expert_idx in range(self.num_experts):
+            for token_idx, output in expert_outputs_dict[expert_idx].items():
+                dense_outputs[token_idx, expert_idx] = output
+        
+        # ---------- 添加伪梯度路径 ----------
+        # 对于每个未激活的专家，创建一个小的梯度路径
+        # 首先创建一个全零的mask标记哪些(token,expert)对已被计算
+        computed_mask = torch.zeros((flat_hidden.size(0), self.num_experts), 
+                                   dtype=torch.bool, device=flat_hidden.device)
+        
+        # 标记已计算的(token,expert)对
+        for expert_idx in range(self.num_experts):
+            for token_idx in token_indices_per_expert[expert_idx]:
+                computed_mask[token_idx, expert_idx] = True
+        
+        # 对未计算的(token,expert)对，添加微小的伪梯度路径
+        # 我们使用一个小的缩放参数来确保前向计算几乎不受影响
+        scale = 1e-4
+        
+        # 计算一个残差连接，确保梯度能够流向所有专家
+        for token_idx in range(flat_hidden.size(0)):
+            token_input = flat_hidden[token_idx:token_idx+1]  # 保持维度
             
-        dense_outputs = torch.cat(dense_outputs, dim=0)  # (B*seq_len, hidden_dim)
-        # ---------- Dense估计结束 ----------
-
-        # 使用直通梯度技巧
-        final_flat = sparse_output.detach() + (dense_outputs - dense_outputs.detach())
+            # 对每个未激活的专家创建伪梯度路径
+            for expert_idx in range(self.num_experts):
+                if not computed_mask[token_idx, expert_idx]:
+                    # 查找是否有token激活了这个专家
+                    if token_indices_per_expert[expert_idx]:
+                        # 从激活该专家的token中选择一个
+                        # 优先选择与当前token共享其他专家的token
+                        similar_tokens = []
+                        for other_idx in token_indices_per_expert[expert_idx]:
+                            # 检查是否有共同激活的专家
+                            common_experts = set(token_to_experts[token_idx]) & set(token_to_experts[other_idx])
+                            if common_experts:
+                                similar_tokens.append(other_idx)
+                        
+                        if similar_tokens:
+                            # 使用相似token的平均输出
+                            similar_outputs = [expert_outputs_dict[expert_idx][t] for t in similar_tokens]
+                            estimated_output = torch.stack(similar_outputs).mean(0)
+                        else:
+                            # 使用所有激活该专家的token的平均输出
+                            all_outputs = [expert_outputs_dict[expert_idx][t] for t in token_indices_per_expert[expert_idx]]
+                            estimated_output = torch.stack(all_outputs).mean(0)
+                            
+                        # 添加微小的直接计算以维持梯度流
+                        direct_output = self.experts[expert_idx](token_input).squeeze(0)
+                        
+                        # 组合估计输出和直接计算
+                        # 前向使用估计输出，反向使用直接计算
+                        combined = estimated_output.detach() + scale * (direct_output - direct_output.detach())
+                        dense_outputs[token_idx, expert_idx] = combined
+                    else:
+                        # 如果没有token激活该专家，直接进行计算但使用小缩放
+                        # 这保证了梯度流而对前向几乎无影响
+                        direct_output = scale * self.experts[expert_idx](token_input).squeeze(0)
+                        dense_outputs[token_idx, expert_idx] = direct_output
+        
+        # ---------- 组合输出 ----------
+        # 使用路由权重对每个token的所有专家输出进行加权
+        dense_combined = torch.zeros_like(sparse_output)
+        
+        for token_idx in range(flat_hidden.size(0)):
+            # 对该token的所有专家输出进行加权
+            token_experts_output = dense_outputs[token_idx]  # (num_experts, hidden_dim)
+            token_routing_weights = routing_weights[token_idx].unsqueeze(-1)  # (num_experts, 1)
+            weighted_experts = token_experts_output * token_routing_weights
+            token_output = weighted_experts.sum(0)  # (hidden_dim)
+            dense_combined[token_idx] = token_output
+        
+        # 使用直通梯度技巧：前向使用sparse_output，反向使用dense_combined
+        # 增大直通梯度系数以增强梯度流
+        straight_through_scale = 1.0  # 可以尝试不同的值
+        final_flat = sparse_output.detach() + straight_through_scale * (dense_combined - dense_combined.detach())
+        
         final_output = final_flat.view(batch_size, seq_length, hidden_dim)
         return final_output, router_logits
 
-    def estimate_dense_output_efficient(self, token_idx, activated, gate_prob, 
-                                       all_activated_outputs, all_routing_indices, token_activated_experts):
-        """
-        优化版本的dense输出估计，只使用已计算的专家输出
-        """
-        num_experts = gate_prob.size(0)
-        dense_parts = {}
-        
-        # 对于激活的专家，直接使用其输出
-        for expert_idx in activated:
-            if expert_idx in all_activated_outputs and token_idx in all_activated_outputs[expert_idx]:
-                dense_parts[expert_idx] = all_activated_outputs[expert_idx][token_idx]
-        
-        # 对于未激活的专家，使用其他token的激活输出估计
-        non_activated = [i for i in range(num_experts) if i not in activated]
-        for expert_idx in non_activated:
-            # 如果该专家没有被任何token激活，跳过
-            if expert_idx not in all_routing_indices or not all_routing_indices[expert_idx]:
-                # 使用零向量或平均值作为估计
-                dense_parts[expert_idx] = torch.zeros_like(next(iter(dense_parts.values()))) if dense_parts else 0
-                continue
-                
-            # 找出激活了该专家的token，并且这些token也激活了当前token激活的某些专家
-            candidate_tokens = []
-            for other_token in all_routing_indices[expert_idx]:
-                # 检查other_token是否与当前token共享某些激活专家
-                if other_token in token_activated_experts:
-                    common_experts = set(activated) & set(token_activated_experts[other_token])
-                    if common_experts:
-                        candidate_tokens.append(other_token)
-            
-            # 如果找到了候选token，使用它们的输出平均值
-            if candidate_tokens:
-                expert_outputs = [all_activated_outputs[expert_idx][t] for t in candidate_tokens]
-                estimated = torch.stack(expert_outputs).mean(dim=0)
-            else:
-                # 找不到合适的候选，使用所有激活了该专家的token
-                expert_outputs = [all_activated_outputs[expert_idx][t] for t in all_routing_indices[expert_idx]]
-                estimated = torch.stack(expert_outputs).mean(dim=0)
-                
-            dense_parts[expert_idx] = estimated
-            
-        # 按路由权重加权求和
-        estimated_dense = 0
-        for expert_idx in range(num_experts):
-            if expert_idx in dense_parts:
-                estimated_dense += gate_prob[expert_idx] * dense_parts[expert_idx]
-                
-        return estimated_dense
-
 
 class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     """