Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -25,32 +25,41 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
     """
     def forward(self, hidden_states: torch.Tensor):
         batch_size, seq_length, hidden_dim = hidden_states.shape
+        # 记录输入张量的数据类型，确保所有计算保持一致
+        dtype = hidden_states.dtype
+        device = hidden_states.device
+        
         flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
         N_tokens = flat_hidden.size(0)
 
         # 计算路由逻辑
         router_logits = self.gate(flat_hidden)  # (B*seq_len, num_experts)
-        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
+        # 确保router_logits和flat_hidden数据类型一致
+        router_logits = router_logits.to(dtype=dtype)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=dtype)  # (B*seq_len, num_experts)
 
         # 选择top-k专家
         routing_weights_topk, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
         if self.norm_topk_prob:
             routing_weights_topk = routing_weights_topk / routing_weights_topk.sum(dim=-1, keepdim=True)
-        routing_weights_topk = routing_weights_topk.to(flat_hidden.dtype)
-
+            # 确保归一化后类型一致
+            routing_weights_topk = routing_weights_topk.to(dtype=dtype)
+        
         # ---------- 真实计算所有专家输出（密集计算）----------
         all_expert_outputs = torch.zeros((N_tokens, self.num_experts, hidden_dim),
-                                         dtype=flat_hidden.dtype, device=flat_hidden.device)
+                                         dtype=dtype, device=device)
         
         for expert_idx in range(self.num_experts):
             expert_layer = self.experts[expert_idx]
             # 对所有token都计算当前专家的输出
             expert_output = expert_layer(flat_hidden)  # (N_tokens, hidden_dim)
+            # 确保专家输出与预期类型一致
+            expert_output = expert_output.to(dtype=dtype)
             all_expert_outputs[:, expert_idx, :] = expert_output
         
         # ---------- 提取激活专家输出（稀疏前向）- 使用张量批处理 ----------
         # 创建索引张量，第一维是token索引，第二维是专家索引
-        token_indices = torch.arange(N_tokens, device=flat_hidden.device).unsqueeze(1).expand(-1, self.top_k)
+        token_indices = torch.arange(N_tokens, device=device).unsqueeze(1).expand(-1, self.top_k)
         batch_indices = token_indices.reshape(-1)
         expert_indices = selected_experts.reshape(-1)
         
@@ -59,6 +68,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         
         # 扩展权重以便批量相乘
         expanded_weights = routing_weights_topk.unsqueeze(-1)  # (N_tokens, top_k, 1)
+        expanded_weights = expanded_weights.to(dtype=dtype)
         
         # 权重乘以专家输出并求和
         sparse_output = (selected_outputs * expanded_weights).sum(dim=1)  # (N_tokens, hidden_dim)
@@ -66,12 +76,14 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         # ---------- 密集计算聚合（用于反向传播）----------
         # 使用所有专家的输出和路由权重计算密集输出
         routing_weights_expanded = routing_weights.unsqueeze(-1)  # (N_tokens, num_experts, 1)
+        routing_weights_expanded = routing_weights_expanded.to(dtype=dtype)
         dense_outputs = (all_expert_outputs * routing_weights_expanded).sum(dim=1)  # (N_tokens, hidden_dim)
         
         # ---------- 组合稀疏前向和密集反向 ----------
         # sparse_output.detach()保留稀疏前向计算图
         # (dense_outputs - dense_outputs.detach())只保留密集反向梯度
         final_flat = sparse_output.detach() + (dense_outputs - dense_outputs.detach())
+        final_flat = final_flat.to(dtype=dtype)  # 确保最终输出类型一致
         final_output = final_flat.view(batch_size, seq_length, hidden_dim)
         
         return final_output, router_logits