diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py index f9da6eb9d350..091ca2b54a24 100644 --- a/deepspeed/moe/sharded_moe.py +++ b/deepspeed/moe/sharded_moe.py @@ -471,13 +471,17 @@ def __init__(self, self.timers = SynchronizedWallClockTimer() self.wall_clock_breakdown = False - self.use_tutel = use_tutel and TUTEL_INSTALLED + self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1 if self.use_tutel: logger.info('Using Tutel optimizations.') elif use_tutel and not TUTEL_INSTALLED: logger.warning("Tutel optimization requested but not installed. " "Proceeding without Tutel.") + elif use_tutel and TUTEL_INSTALLED and gate.k != 1: + logger.warning( + "To enable Tutel optimization, use top-1 instead of top-2 gate. " + "Proceeding without Tutel.") def _set_ep_group(self, ep_group): self.ep_group = ep_group