自定义文件

__init__.py

@@ -1,6 +1,6 @@
 # 导出自定义配置和模型类
-from .configuration_densebackward_olmoe0125 import DenseBackwardOLMoEConfig
-from .modeling_densebackward_olmoe0125 import DenseBackwardOLMoEForCausalLM, DenseBackwardOlmoeSparseMoeBlock
+from .configuration_densebackward_olmoe0125_v1 import DenseBackwardOLMoEConfig
+from .modeling_densebackward_olmoe0125_v1 import DenseBackwardOLMoEForCausalLM, DenseBackwardOlmoeSparseMoeBlock
 
 # 显式注册模型类型
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING

config.json

@@ -4,9 +4,9 @@
     "DenseBackwardOLMoEForCausalLM"
   ],
   "auto_map": {
-  "AutoConfig": "configuration_densebackward_olmoe0125.DenseBackwardOLMoEConfig",
-  "AutoModel": "modeling_densebackward_olmoe0125.DenseBackwardOLMoEForCausalLM",
-  "AutoModelForCausalLM": "modeling_densebackward_olmoe0125.DenseBackwardOLMoEForCausalLM"
+  "AutoConfig": "configuration_densebackward_olmoe0125_v1.DenseBackwardOLMoEConfig",
+  "AutoModel": "modeling_densebackward_olmoe0125_v1.DenseBackwardOLMoEForCausalLM",
+  "AutoModelForCausalLM": "modeling_densebackward_olmoe0125_v1.DenseBackwardOLMoEForCausalLM"
 },
   "attention_bias": false,
   "attention_dropout": 0.0,

modeling_densebackward_olmoe0125.py → modeling_densebackward_olmoe0125_v1.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 
 # 导入官方实现（注意根据你的 transformers 版本调整导入路径）
 from transformers.models.olmoe.modeling_olmoe import OlmoeForCausalLM, OlmoeSparseMoeBlock, OlmoeMLP
-from .configuration_densebackward_olmoe0125 import DenseBackwardOLMoEConfig
+from .configuration_densebackward_olmoe0125_v1 import DenseBackwardOLMoEConfig
 
 
 class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
@@ -77,6 +77,10 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         # 使用所有专家的输出和路由权重计算密集输出
         routing_weights_expanded = routing_weights.unsqueeze(-1)  # (N_tokens, num_experts, 1)
         routing_weights_expanded = routing_weights_expanded.to(dtype=dtype)
+        print(expanded_weights.shape)
+        print("sparse",expanded_weights)
+        print(routing_weights_expanded.shape)
+        print("dense",routing_weights_expanded)
         dense_outputs = (all_expert_outputs * routing_weights_expanded).sum(dim=1)  # (N_tokens, hidden_dim)
         
         # ---------- 组合稀疏前向和密集反向 ----------