chg: disable moe amax routing remedy when router group is not None.

chg: remove global barrier in SequentialMLP
Signed-off-by: Chenhan Yu <[email protected]>

Author: ChenhanYu

modelopt/torch/quantization/plugins/megatron.py

@@ -581,7 +581,7 @@ def sync_moe_local_experts_amax(self):
         This function is called to synchronize the amax values across local experts s.t. all localexperts will
         share the same amax.
         """
-        torch.distributed.barrier()
+        # torch.distributed.barrier()
         # Collect amax from all local experts
         amax_dict = {}
         for expert in self.local_experts:
@@ -754,8 +754,11 @@ def _setup(self):
 
     def forward(self, hidden_states):
         if any(getattr(m, "_if_calib", False) for m in self.experts.modules()):
-            original_top_k = self.router.topk
-            self.router.topk = self.router.num_experts
-            super().forward(hidden_states)
-            self.router.topk = original_top_k
+            if self.config.moe_router_num_groups is None:
+                original_top_k = self.router.topk
+                self.router.topk = self.router.num_experts
+                super().forward(hidden_states)
+                self.router.topk = original_top_k
+            else:
+                super().forward(hidden_states)
         return super().forward(hidden_states)