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).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:
-                estimated = all_expert_outputs[:, i, :].mean(dim=0)
+                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
@@ -152,20 +170,48 @@ class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     base_model_prefix = "olmoe"
 
     def __init__(self, config):
+        # 首先调用父类初始化方法
         super().__init__(config)
+        
+        # 不要尝试重新赋值self，而是从预训练模型加载并更新当前模型
+        pretrained_model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0125")
+        
+        # 复制预训练模型的状态到当前模型
+        self.config = pretrained_model.config
+        self.model = pretrained_model.model
+        self.vocab_size = pretrained_model.vocab_size
+        self.router_aux_loss_coef = pretrained_model.router_aux_loss_coef
+        self.num_experts = pretrained_model.num_experts
+        self.lm_head = pretrained_model.lm_head
+        
         # 遍历模型中所有 decoder 层，替换每个 OlmoeSparseMoeBlock 为 DenseBackward 版本
         # 此处假设官方模型在 self.model.layers 中组织 decoder 层，
         # 且每层中 mlp 模块包含属性 sparse_moe_block。
         for layer in self.model.layers:
-            if hasattr(layer.mlp, "sparse_moe_block"):
-                orig_block = layer.mlp.sparse_moe_block
+            if hasattr(layer.mlp, "gate"):
+                print("111")
+                orig_block = layer.mlp
                 # 通过直接复制原版属性创建新的块
                 new_block = DenseBackwardOlmoeSparseMoeBlock(config)  # 或其他适当参数
                 # 然后手动复制需要共享的属性：
                 new_block.gate = orig_block.gate
                 new_block.experts = orig_block.experts
-                new_block.router = orig_block.router
                 new_block.num_experts = orig_block.num_experts
                 new_block.top_k = orig_block.top_k
                 new_block.norm_topk_prob = orig_block.norm_topk_prob
-                layer.mlp.sparse_moe_block = new_block
+                layer.mlp = new_block
+                print(type(layer.mlp))
+
+def main():
+    config = DenseBackwardOLMoEConfig(        # 官方模型参数
+    model_marker="DenseBackward_olmoe_marker",
+)
+# 创建自定义模型实例
+    model = DenseBackwardOLMoEForCausalLM(config)
+    print(type(model))
+    print(type(model.model))
+    print(type(model.model.layers[0]))
+    print(type(model.model.layers[0].mlp))
+    print(type(model.model.layers[0].mlp.experts))
+if __name__ == "__main__":
+    main()