fix(test): resolve flakiness on moe_tests.py

[mesakhcienet]

Description

This PR resolves persistent flakiness in the tests/unit/moe_test.py suite observed specifically in the Nightly CI pipelines. It addresses the root cause of the numerical drift by implementing a stable float32 reference comparison.

Problem Context

Recent changes in the bleeding-edge XLA compiler (pulled in the nightly pipeline) introduced slight variations in the accumulation order for bfloat16 matrix multiplications.
Because our original tests were running both the reference implementation (get_expected_output) and the production kernel (get_moe_output) in bfloat16, we were doing a "noise-vs-noise" comparison. The combined compiler drift on both sides of the equation pushed the accumulation error to ~0.015, exceeding our strict 0.01 test tolerance.

What We Tried & Why We Chose This Solution

1. Attempt 1: Loosening Tolerances (The "Bad" Fix)
   Initially (in a reverted commit), we tried simply raising the rtol and atol from 1e-2 to 5e-2.
   • Why we abandoned this: A 5e-2 tolerance is massive and completely masks real logic regressions, such as MoE routing collapse bugs (e.g., returning all NaNs, which passed silently because equal_nan=False was dropped).
2. Attempt 2: The float32 Reference (The Correct Fix)
   We changed the reference MoeLoopBlock to compute entirely in jnp.float32, while keeping the actual production code in bfloat16.
   • Why this is the best practice: This creates a "signal-vs-noise" test. By pinning the reference to an exact mathematical truth, we eliminate half the variance in the test. The maximum deviation against the nightly compiler now measures reliably at ~0.0058, safely passing the strict 0.01 threshold without needing to weaken our test bounds.

Implementation Details

• Upcasted Reference: get_expected_output is now initialized with dtype=jnp.float32 and hidden_states are upcasted before the forward pass.
• Pinned Gate Logits: We added float32_gate_logits=True to the pyconfig.initialize blocks. This ensures that the float32 reference and the bfloat16 production code make identical Top-K routing decisions, preventing massive divergence caused by floating-point rounding at the routing boundary.
• Introduced assert_moe_close: We removed the opaque self.assertTrue(jax.numpy.allclose(...)) wrappers and introduced a custom helper. This runs np.testing.assert_allclose to provide rich diffs, and explicitly prints the max_diff and max_diff/RMS before evaluating. This drastically improves observability for future compiler changes.

BUGS: https://buganizer.corp.google.com/issues/536702731

Tests

To prove the fix works, we reproduced the failing CI environment locally:

1. Ran bash src/dependencies/scripts/setup.sh MODE=nightly to pull the bleeding-edge XLA/libtpu dependencies that trigger the math drift.
2. Ran python3 -m pytest tests/unit/moe_test.py on the main branch to observe the legitimate failures.
3. Applied the float32 reference fix and reran the suite.
4. Verified that the max_diff hovered at 0.0058 and all tests passed perfectly within the original strict rtol=2e-2, atol=1e-2 bounds.

We also run the chages on the github workflow runner in here

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!