Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -80,7 +80,13 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
 
         # ---------- Dense估计部分 ----------
         # 计算所有专家对所有 token 的 dense 输出，shape: (B*seq_len, num_experts, hidden_dim)
-        all_expert_outputs = torch.stack([expert(flat_hidden) for expert in self.experts], dim=1)
+        # 创建全零张量，只填入已激活专家的输出
+        all_expert_outputs = torch.zeros((flat_hidden.size(0), self.num_experts, hidden_dim), 
+                                   dtype=flat_hidden.dtype, device=flat_hidden.device)
+        # 填入已激活专家的输出
+        for i in range(flat_hidden.size(0)):
+            for expert_idx in activated_outputs[i].keys():
+                all_expert_outputs[i, expert_idx] = activated_outputs[i][expert_idx]
         # 将 selected_experts 转换为 list，每个 token 的激活专家列表
         all_routing = selected_experts.tolist()  # 长度为 (B*seq_len)
         
@@ -130,9 +136,21 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
                     indices.append(idx)
             if indices:
                 selected_outputs = all_expert_outputs[indices, i, :]  # (n, hidden_dim)
-                estimated = selected_outputs.mean(dim=0)
+                 # 只计算非零值的平均值
+                mask = (selected_outputs.sum(dim=-1) != 0).float().unsqueeze(-1)  # (n, 1)
+                if mask.sum() > 0:
+                    estimated = (selected_outputs * mask).sum(dim=0) / mask.sum()
+                else:
+                   # 如果全是零，返回零向量
+                    estimated = torch.zeros_like(selected_outputs[0])
             else:
-                estimated = all_expert_outputs[:, i, :].mean(dim=0)
+                all_outputs = all_expert_outputs[:, i, :]
+                mask = (all_outputs.sum(dim=-1) != 0).float().unsqueeze(-1)  # (N, 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