Skip to content

Guard zero-sum MoE gate normalization - #1375

Open
WangXukang-cypher wants to merge 1 commit into
apple:mainfrom
WangXukang-cypher:fix/moe-epsilon-guard
Open

WangXukang-cypher wants to merge 1 commit into
apple:mainfrom
WangXukang-cypher:fix/moe-epsilon-guard

Conversation

@WangXukang-cypher

Copy link
Copy Markdown

Summary

This PR adds epsilon guards to MoE gate normalization paths to avoid NaNs when the selected gate weights sum to zero.

Specifically, it updates the normalization logic in:

  • TopKGating
  • TopKDropFreeGating
  • TopKBiasGating

It also adds regression tests covering zero-score / zero-sum gate cases, including the bias-gating group-routing path.

Test Plan

  • Ran 27 existing MoE tests: all passed
  • Observed 1 failing test: test_moe_quantized_einsum0
    • This is a pre-existing TPU-only quantized test that fails on CPU
    • It is unrelated to this epsilon fix
  • Ran targeted regression tests for this change in the local conda axlearn environment: 5 passed
    • TopKGatingTest::test_sigmoid_score_fn
    • TopKGatingTest::test_zero_score_fn_produces_finite_combine_tensor
    • TopKDropFreeGatingTest
    • TopKBiasGatingTest

@WangXukang-cypher
WangXukang-cypher requested a review from a team as a code owner September 4, 2026 06:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant