fix moe gate dtype, fix tbo, fix fake dispatch (#7825)

2025-07-19 22:13:46 -07:00
parent f62d75b6a1
commit a589a07167
3 changed files with 3 additions and 3 deletions
@@ -229,7 +229,7 @@ class MoEGate(nn.Module):
        )
        if config.topk_method == "noaux_tc":
            self.e_score_correction_bias = nn.Parameter(
-                torch.empty((config.n_routed_experts))
+                torch.empty((config.n_routed_experts), dtype=torch.float32)
            )
        else:
            self.e_score_correction_bias = None