Skip to content

Fix moe amax remedy for dsr1 and remove global barrier in quantization megatron plugins #808

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChenhanYu merged 2 commits into from
Jan 23, 2026
Merged

Fix moe amax remedy for dsr1 and remove global barrier in quantization megatron plugins #808

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5bc77c5
chg: disable moe amax routing remedy when router group is not None.
ChenhanYu Jan 16, 2026
606cfba
chg: remove comments
ChenhanYu Jan 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions modelopt/torch/quantization/plugins/megatron.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,7 +581,6 @@ 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()
# Collect amax from all local experts
amax_dict = {}
for expert in self.local_experts:
Expand Down Expand Up @@ -754,8 +753,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)