Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -137,7 +137,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
             if indices:
                 selected_outputs = all_expert_outputs[indices, i, :]  # (n, hidden_dim)
                  # 只计算非零值的平均值
-                mask = (selected_outputs.sum(dim=-1) != 0).float().unsqueeze(-1)  # (n, 1)
+                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:
@@ -145,7 +145,7 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
                     estimated = torch.zeros_like(selected_outputs[0])
             else:
                 all_outputs = all_expert_outputs[:, i, :]
-                mask = (all_outputs.sum(dim=-1) != 0).float().unsqueeze(-1)  # (N, 1)
+                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: