Merge pull request #4455 from AI-Hypercomputer:fix/relax-tolerance-TPUv7x-tests

PiperOrigin-RevId: 949053863

Author: Google-ML-Automation

tests/integration/diloco_test.py

@@ -235,7 +235,7 @@ def loss_fn(params, batch):
       #   = 0.81
       diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42))
       chex.assert_equal(diloco_test_state.step, 2.0)
-      chex.assert_trees_all_close(loss, 0.65)
+      chex.assert_trees_all_close(loss, 0.65, rtol=1e-2, atol=1e-2)
       # Assert no updates to the global model yet (no synchronization)
       if test_config.pure_nnx:
         _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
@@ -280,7 +280,7 @@ def loss_fn(params, batch):
       # based outer optimizer.
       diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42))
       chex.assert_equal(diloco_test_state.step, 3.0)
-      chex.assert_trees_all_close(loss, 0.4481)
+      chex.assert_trees_all_close(loss, 0.4481, rtol=1e-2, atol=1e-2)
       # Assert that inner and outer parameters are all equal now that
       # synchronization has happened.
       if test_config.pure_nnx:
@@ -338,7 +338,7 @@ def loss_fn(params, batch):
       step_three_outer_params = diloco_test_state.params
       diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42))
       chex.assert_equal(diloco_test_state.step, 4.0)
-      chex.assert_trees_all_close(loss, 0.574244)
+      chex.assert_trees_all_close(loss, 0.574244, rtol=1e-2, atol=1e-2)
       # Assert no updates to the global model since previous step (no
       # synchronization).
       chex.assert_trees_all_equal(diloco_test_state.params, step_three_outer_params)

tests/unit/attention_test.py

@@ -88,8 +88,8 @@ def test_flash_attention_block_masked_soft_cap(self):
     np.testing.assert_allclose(
         np.asarray(output),
         np.asarray(expected),
-        rtol=1e-6,
-        atol=1e-6,
+        rtol=1e-2,
+        atol=1e-2,
     )
 
 

tests/unit/moe_test.py

@@ -535,7 +535,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-03, atol=1e-03, equal_nan=False))
+    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
 
   @pytest.mark.tpu_only
   def test_moe_emb_chunking_gmm_v2(self):