Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -152,21 +152,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()