Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -66,23 +66,12 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
                 weighted_output = current_output * weight
                 sparse_output.index_add_(0, top_x, weighted_output.to(flat_hidden.dtype))
                 
-                # 保存专家输出到张量中，而不是使用字典
-                all_expert_outputs.index_copy_(0, top_x, 
-                                             torch.zeros_like(all_expert_outputs[0:top_x.size(0)]).scatter_(
-                                                 1, expert_idx * torch.ones((top_x.size(0), 1), 
-                                                                         dtype=torch.long, 
-                                                                         device=flat_hidden.device), 
-                                                 current_output.unsqueeze(1)))
+                # 直接为激活的token分配专家输出
+                for i in range(top_x.shape[0]):
+                    token_idx = top_x[i]
+                    all_expert_outputs[token_idx, expert_idx] = current_output[i]
+                    expert_activated[token_idx, expert_idx] = True
                 
-                # 标记哪些专家被激活
-                expert_activated.index_copy_(0, top_x,
-                                           torch.zeros_like(expert_activated[0:top_x.size(0)]).scatter_(
-                                               1, expert_idx * torch.ones((top_x.size(0), 1), 
-                                                                       dtype=torch.long, 
-                                                                       device=flat_hidden.device), 
-                                               torch.ones((top_x.size(0), 1), 
-                                                        dtype=torch.bool, 
-                                                        device=flat_hidden.device)))
         # ---------- 稀疏计算结束 ----------
 
         # ---------- Dense估计部分 ----------
@@ -126,7 +115,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         num_experts = routing_weights.size(1)
         device = all_expert_outputs.device
         
-        # 预分配结果张量
+        # 预分配结果张量，注意是hidden_dim而不是num_experts
         dense_outputs = torch.zeros((total_tokens, hidden_dim), dtype=all_expert_outputs.dtype, device=device)
         
         # 对每个token单独处理（此处仍需循环，但后续可进一步优化）
@@ -191,6 +180,55 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         
         return dense_outputs
 
+    def estimate_dense_output(self, token_idx, activated, gate_prob, activated_outputs, all_routing, all_expert_outputs):
+        """
+        对于当前 token，根据 mini-batch 中的信息估计 dense 输出。
+        参数：
+            token_idx: 当前 token 的索引（标量）
+            activated: 当前 token 激活的专家列表，例如 [1, 3]
+            gate_prob: 当前 token 的 routing 权重，形状 (num_experts,)
+            activated_outputs: dict，当前 token 对激活专家的实际输出，形状 (hidden_dim,)
+            all_routing: list，每个 token 的激活专家列表（长度为 N，每个元素为 list）
+            all_expert_outputs: Tensor, (N, num_experts, hidden_dim)
+        返回：
+            estimated_dense: Tensor, (hidden_dim,)
+        """
+        num_experts = gate_prob.size(0)
+        dense_parts = {}
+        # 对于激活的专家，直接使用其实际输出
+        for idx in activated:
+            dense_parts[idx] = activated_outputs[idx]
+        # 对于未激活的专家，使用 mini-batch 中其他 token 的输出估计
+        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
+        # 按 gate_prob 加权求和各专家输出
+        estimated_dense = 0
+        for i in range(num_experts):
+            estimated_dense += gate_prob[i] * dense_parts[i]
+        return estimated_dense
+
 
 class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     """