[PP] Fix disabled flaky tests (#154856)

H-Huang · pytorchmergebot · commit 283f876ab6c2 · 2025-06-03T15:55:29.000Z
Fix #154373, #154391, #154408, #154443, #154481 Because MultiProcContinousTest [now executes the tests with 8 GPUs instead of 2](#153653), our PP tests comparing gradients have become flakier due to the longer pipeline. The gradients are still close but we need to relax the tolerance. Pull Request resolved: #154856 Approved by: https://github.com/Skylion007
diff --git a/test/distributed/pipelining/test_schedule_multiproc.py b/test/distributed/pipelining/test_schedule_multiproc.py
@@ -513,7 +513,7 @@ def test_grad_with_manual_interleaved(self, ScheduleClass, use_new_runtime):
             for name, p in stage_module.named_parameters():
                 ref_p = ref_submod.get_parameter(name)
                 try:
-                    torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
+                    torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=1e-3)
                 except AssertionError:
                     print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
                     raise
