[WangXukang-cypher]

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

[dajiaohuang]

I checked the denominator arithmetic in a CPU-only synthetic JAX 0.8.3 case (not the full MoE forward). With the supported sigmoid scoring function and logits [-25, -25], the current normalization yields [0.5, 0.5], while adding 1e-9 yields about [0.0135, 0.0135] even though the denominator is nonzero. That changes valid small positive gates, not only the zero-sum case. Could the guard replace only a zero denominator (for example, safe_denom = jnp.where(denom == 0, 1, denom)) and add a small-positive regression alongside the all-zero case? The same epsilon pattern appears at all five normalization sites in this diff.

[WangXukang-cypher]

Thanks for catching this and providing a concrete example! You're right: adding 1e-9 changes the normalization for small positive gate sums, which goes beyond the intended zero-sum fix.

I'll update all five normalization sites to guard only zero denominators using jnp.where(denom == 0, 1, denom). I'll also add regression tests with sigmoid scoring and small positive gate sums, including the [-25, -25] case, while keeping the existing all-zero tests.