Skip to content

Fix modality MoE padding routing and selective recompute - #7763

Open
HollowMan6 wants to merge 3 commits into
NVIDIA:pull-request/7502from
HollowMan6:modality-moe-padding-bias-fixes
Open

HollowMan6 wants to merge 3 commits into
NVIDIA:pull-request/7502from
HollowMan6:modality-moe-padding-bias-fixes

Conversation

@HollowMan6

@HollowMan6 HollowMan6 commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #7502.

Summary

  • Keep dispatcher rows for padded tokens, zero their routing probabilities, and exclude them from load-balancing counters.
  • Reuse the regular MoE checkpoint path for selective recompute in the modality MoE layer.
  • Skip expert-bias collectives when the update rate is zero.
  • Add unit tests in tests/unit_tests/models/test_deepseek_v41_components.py.

@copy-pr-bot

copy-pr-bot Bot commented Sep 30, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Preserve dispatcher rows for padded tokens while zeroing their probabilities and excluding them from balance counters. Reuse the regular MoE checkpoint path and avoid unnecessary expert-bias collectives at zero update rate.

Signed-off-by: Hollow Man <hollowman@opensuse.org>
@HollowMan6
HollowMan6 force-pushed the modality-moe-padding-bias-fixes branch from b16c8ab to cecf595 Compare September 30, 2026 22:32
@HollowMan6

Copy link
Copy Markdown
Contributor Author

@guihong-nv This is a follow-up PR on top of #7502. It fixes padded-token routing by keeping their dispatcher rows while zeroing their probs and leaving them out of the balance counters. It also reuses the regular MoE checkpoint path for selective recompute and skips expert-bias collectives when the update rate is 0. Verified together with NVIDIA-NeMo/Megatron-Bridge#6199. Please let me know if you have any comments.

@HollowMan6
HollowMan6 marked this pull request as ready for review September 30, 2026 22:37
@HollowMan6
HollowMan6 requested review from a team as code owners September 30, 2026 22:37
Copilot AI balanced review requested due to automatic review settings September 30, 2026 22:37

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Oct 3, 2026
@dimapihtar dimapihtar added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Oct 5, 2026
@dimapihtar

Copy link
Copy Markdown
Contributor

/ok to test cecf595

@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Oct 5, 2026
Signed-off-by: Hollow Man <hollowman@opensuse.org>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Oct 7, 2026
Signed-off-by: Hollow Man <hollowman@opensuse.org>

This branch was successfully deployed

1 active (outdated) deployment
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants