Merge pull request #4518 from AI-Hypercomputer:fix/routed-moe-test

PiperOrigin-RevId: 950968672

Author: Google-ML-Automation

tests/unit/moe_test.py

@@ -29,12 +29,31 @@
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import NdInitializer, nd_dense_init, variable_to_logically_partitioned
 from maxtext.layers.quantizations import Fp8Quantization
-from maxtext.utils import maxtext_utils
+from maxtext.utils import max_logging, maxtext_utils
 from maxtext.utils.sharding import remove_expert_from_partition_spec
 from tests.utils.test_helpers import get_test_config_path
 import pytest
 
 
+def assert_moe_close(actual, expected, dtype):
+  """Asserts that the actual and expected MoE outputs are close."""
+  assert np.isfinite(actual).all(), "Actual output contains NaNs or Infs!"
+  if dtype == jnp.bfloat16:
+    rtol, atol = 2e-2, 1e-2
+  else:
+    rtol, atol = 1e-5, 1e-6
+
+  max_diff = float(np.max(np.abs(actual - expected)))
+  rms_expected = float(np.sqrt(np.mean(np.square(expected))))
+  max_logging.debug(
+      f"\n[assert_moe_close] dtype={dtype}, max_diff={max_diff:.6f}, max_diff/RMS={max_diff/rms_expected:.6f}"
+  )
+
+  np.testing.assert_allclose(
+      np.array(actual, dtype=np.float32), np.array(expected, dtype=np.float32), rtol=rtol, atol=atol, equal_nan=False
+  )
+
+
 class TokenDroppingTest(unittest.TestCase):
 
   def setUp(self):
@@ -156,7 +175,7 @@ def test_generate_masks(self):
     actual_dispatch_mask, actual_combine_mask = self.model.generate_masks(top_k_indices, softmax_probs)
 
     self.assertTrue((expected_dispatch_mask == actual_dispatch_mask).all())
-    self.assertTrue(jax.numpy.allclose(expected_combine_mask, actual_combine_mask, rtol=1e-02, atol=1e-02))
+    assert_moe_close(actual_combine_mask, expected_combine_mask, self.cfg.dtype)
 
 
 class MlpBlockTest(unittest.TestCase):
@@ -280,7 +299,7 @@ def test_deepseek_bias_updates(self):
     expected_updates = jnp.array([-0.01, 0.0, 0.01, 0.0])
     actual_updates = moe.calculate_load_balance_updates(top_k_indices, num_experts, rate)
 
-    self.assertTrue(jax.numpy.allclose(expected_updates, actual_updates, rtol=1e-05, atol=1e-05, equal_nan=False))
+    assert_moe_close(actual_updates, expected_updates, jnp.float32)
 
 
 class MoeLoopBlock(nnx.Module):
@@ -394,7 +413,7 @@ def get_expected_output(self, rng, hidden_states, cfg, mesh):
         num_experts_per_tok=cfg.num_experts_per_tok,
         kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
         kernel_axes=("embed", "mlp"),
-        dtype=cfg.dtype,
+        dtype=jnp.float32,
         weight_dtype=cfg.weight_dtype,
     )
     variables = model.init(
@@ -404,8 +423,8 @@ def get_expected_output(self, rng, hidden_states, cfg, mesh):
         ),
     )
 
-    output = jax.jit(model.apply)(variables, hidden_states)  # pylint: disable=not-callable
-    return variables, output
+    output = jax.jit(model.apply)(variables, hidden_states.astype(jnp.float32))  # pylint: disable=not-callable
+    return variables, output.astype(cfg.dtype)
 
   def get_moe_output(self, variables, hidden_states, cfg, mesh):
     """retrieve expected output from MoE"""
@@ -462,6 +481,7 @@ def test_megablox(self):
         sparse_matmul=True,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(1234)
@@ -477,7 +497,7 @@ def test_megablox(self):
     mesh = Mesh(devices_array, cfg.mesh_axes)
     variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
     actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+    assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_ragged_dot(self):
@@ -491,6 +511,7 @@ def test_ragged_dot(self):
         sparse_matmul=True,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(1234)
@@ -506,7 +527,7 @@ def test_ragged_dot(self):
     mesh = Mesh(devices_array, cfg.mesh_axes)
     variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
     actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+    assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_dense(self):
@@ -520,6 +541,7 @@ def test_dense(self):
         sparse_matmul=False,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -535,7 +557,7 @@ def test_dense(self):
     mesh = Mesh(devices_array, cfg.mesh_axes)
     variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
     actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+    assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_moe_emb_chunking_random_routing(self):
@@ -556,6 +578,7 @@ def test_moe_emb_chunking_random_routing(self):
         mlp_bias=True,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     cfg_non_chunked = pyconfig.initialize(
@@ -575,6 +598,7 @@ def test_moe_emb_chunking_random_routing(self):
         mlp_bias=True,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(1234)
@@ -639,6 +663,7 @@ def test_moe_emb_chunking_gmm_v2(self):
         mlp_bias=True,
         per_device_batch_size=1,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(1234)
@@ -654,7 +679,7 @@ def test_moe_emb_chunking_gmm_v2(self):
     mesh = Mesh(devices_array, cfg.mesh_axes)
     variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
     actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+    assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_megablox_expert_parallelism(self):
@@ -669,6 +694,7 @@ def test_megablox_expert_parallelism(self):
         per_device_batch_size=4,  # TODO(b/450900273): sharding error if pdbs=1
         ici_expert_parallelism=4,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -685,7 +711,7 @@ def test_megablox_expert_parallelism(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_ring_of_expert_and_tensor_parallelism(self):
@@ -702,6 +728,7 @@ def test_ring_of_expert_and_tensor_parallelism(self):
         use_ring_of_experts=True,
         ici_tensor_parallelism=2,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -718,7 +745,7 @@ def test_ring_of_expert_and_tensor_parallelism(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   def _run_ragged_sort_loss_and_grad(
       self,
@@ -754,6 +781,7 @@ def _build_cfg(use_ragged_sort: bool):
           ici_expert_parallelism=2,
           use_ring_of_experts=use_ring_of_experts,
           max_target_length=128,
+          float32_gate_logits=True,
           use_ragged_sort=use_ragged_sort,
           ragged_buffer_factor=effective_buffer_factor,
           ragged_gather_fallback=ragged_gather_fallback,
@@ -890,6 +918,7 @@ def test_moe_fsdp_two_stage_parallelism_tpu_only(self):
         ici_fsdp_transpose_parallelism=2,
         moe_fsdp_use_two_stage_all_gather=True,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -906,7 +935,7 @@ def test_moe_fsdp_two_stage_parallelism_tpu_only(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_megablox_context_parallelism(self):
@@ -921,6 +950,7 @@ def test_megablox_context_parallelism(self):
         per_device_batch_size=1,
         ici_context_parallelism=4,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -937,7 +967,7 @@ def test_megablox_context_parallelism(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_megablox_expert_context_parallelism(self):
@@ -954,6 +984,7 @@ def test_megablox_expert_context_parallelism(self):
         ici_expert_parallelism=2,
         packing=False,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -970,7 +1001,7 @@ def test_megablox_expert_context_parallelism(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   @pytest.mark.tpu_only
   def test_megablox_expert_tensor_parallelism(self):
@@ -986,6 +1017,7 @@ def test_megablox_expert_tensor_parallelism(self):
         ici_tensor_parallelism=2,
         ici_expert_parallelism=2,
         max_target_length=128,
+        float32_gate_logits=True,
     )
 
     rng = jax.random.PRNGKey(2345)
@@ -1002,7 +1034,7 @@ def test_megablox_expert_tensor_parallelism(self):
     with nn_partitioning.axis_rules(cfg.logical_axis_rules):
       variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
       actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
-      self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
+      assert_moe_close(actual_output, expected_output, cfg.dtype)
 
   def test_random_routing(self):
     bs, seq_len, num_experts, num_experts_per_tok = 12, 1024, 8, 2