Relax tolerances in DiLoCo test to fix TPU7x failure

[hsuan-lun-chiang]

Summary

Relax tolerances in DiLoCoTest to fix TPU7x integration test failure caused by numerical divergence.

Root Cause Analysis

The integration test tests/integration/diloco_test.py::DiLoCoTest::test_diloco_training_simulation_with_mesh was failing on TPU7x CI. Investigation showed that TPU7x exhibits a small numerical divergence (~0.38% relative difference in loss: 0.647502 actual vs 0.65 expected) during DiLoCo training simulation. The default tolerance of assert_trees_all_close (1e-6) was too strict for this hardware-specific divergence, leading to test failures. TPUv6e did not exhibit this divergence (yielding 0.65000009).

Resolution

Relaxed the tolerance parameters (rtol=1e-2, atol=1e-2) for all three assert_trees_all_close checks in diloco_test.py. This allows the test to pass on TPU7x while still verifying the correctness of the training simulation.

Verification

1. Local Verification: Ran the test locally on TPUv6e to ensure it still passes (loss: 0.65000009).
2. CI Verification: Temporarily modified the CI workflow to run TPU7x tests on the PR branch. Verified that tpu7x-integration job successfully passed with the relaxed tolerances.
3. Workflow Restored: Restored the CI workflow to skip TPU7x tests on PRs to conserve resources.

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!