Remove flakiness of test_kv_zch_load_state_dict (#3292)

Summary:
Pull Request resolved: #3292

Test link: https://www.internalfb.com/intern/test/281475203207916

The test is flaky because KVZCH kernel [guarantees accuracy of 1e-2](https://www.internalfb.com/code/fbsource/[35a43c0e43e5]/fbcode/deeplearning/fbgemm/fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py?lines=1399-1402) for FP16.
I changed test_model_parallel_base to accept custom tolerance to override default atol/rtol and added the tolerance to this test to resolve the flakiness

Reviewed By: duduyi2013

Differential Revision: D80457783

fbshipit-source-id: 07720dfceb5a2d393bff2fa2e4e0b0f81c7cac6e

Author: kausv

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -380,26 +380,31 @@ def _eval_models(
         m2: DistributedModelParallel,
         batch: ModelInput,
         is_deterministic: bool = True,
+        tolerance: Optional[float] = None,
     ) -> None:
         with torch.no_grad():
             loss1, pred1 = m1(batch)
             loss2, pred2 = m2(batch)
-
         if is_deterministic:
             self.assertTrue(torch.equal(loss1, loss2))
             self.assertTrue(torch.equal(pred1, pred2))
         else:
-            rtol, atol = _get_default_rtol_and_atol(loss1, loss2)
-            torch.testing.assert_close(loss1, loss2, rtol=rtol, atol=atol)
-            rtol, atol = _get_default_rtol_and_atol(pred1, pred2)
-            torch.testing.assert_close(pred1, pred2, rtol=rtol, atol=atol)
+            if tolerance:
+                torch.testing.assert_close(loss1, loss2, rtol=tolerance, atol=tolerance)
+                torch.testing.assert_close(pred1, pred2, rtol=tolerance, atol=tolerance)
+            else:
+                rtol, atol = _get_default_rtol_and_atol(loss1, loss2)
+                torch.testing.assert_close(loss1, loss2, rtol=rtol, atol=atol)
+                rtol, atol = _get_default_rtol_and_atol(pred1, pred2)
+                torch.testing.assert_close(pred1, pred2, rtol=rtol, atol=atol)
 
     def _compare_models(
         self,
         m1: DistributedModelParallel,
         m2: DistributedModelParallel,
         is_deterministic: bool = True,
         use_virtual_table: bool = False,
+        tolerance: Optional[float] = None,
     ) -> None:
         sd1 = m1.state_dict()
         sd2 = m2.state_dict()
@@ -437,7 +442,12 @@ def _compare_models(
                     if is_deterministic:
                         self.assertTrue(torch.allclose(src_tensor, dst_tensor))
                     else:
-                        rtol, atol = _get_default_rtol_and_atol(src_tensor, dst_tensor)
+                        if tolerance:
+                            rtol, atol = tolerance, tolerance
+                        else:
+                            rtol, atol = _get_default_rtol_and_atol(
+                                src_tensor, dst_tensor
+                            )
                         torch.testing.assert_close(
                             src_tensor, dst_tensor, rtol=rtol, atol=atol
                         )
@@ -453,7 +463,10 @@ def _compare_models(
                     if is_deterministic:
                         self.assertTrue(torch.equal(src, dst))
                     else:
-                        rtol, atol = _get_default_rtol_and_atol(src, dst)
+                        if tolerance:
+                            rtol, atol = tolerance, tolerance
+                        else:
+                            rtol, atol = _get_default_rtol_and_atol(src, dst)
                         torch.testing.assert_close(
                             src._local_tensor, dst._local_tensor, rtol=rtol, atol=atol
                         )
@@ -463,7 +476,10 @@ def _compare_models(
                 if is_deterministic:
                     self.assertTrue(torch.equal(src, dst))
                 else:
-                    rtol, atol = _get_default_rtol_and_atol(src, dst)
+                    if tolerance:
+                        rtol, atol = tolerance, tolerance
+                    else:
+                        rtol, atol = _get_default_rtol_and_atol(src, dst)
                     torch.testing.assert_close(src, dst, rtol=rtol, atol=atol)
 
 

torchrec/distributed/tests/test_model_parallel_nccl_ssd_single_gpu.py

@@ -1022,7 +1022,7 @@ def test_kv_zch_load_state_dict(
             "learning_rate": 0.1,
             "stochastic_rounding": stochastic_rounding,
         }
-        is_deterministic = dtype == DataType.FP32 or not stochastic_rounding
+        is_deterministic = dtype == DataType.FP32
         constraints = {
             table.name: ParameterConstraints(
                 sharding_types=[sharding_type],
@@ -1049,9 +1049,15 @@ def test_kv_zch_load_state_dict(
 
         if is_training:
             self._train_models(m1, m2, batch)
-        self._eval_models(m1, m2, batch, is_deterministic=is_deterministic)
+        self._eval_models(
+            m1, m2, batch, is_deterministic=is_deterministic, tolerance=1e-2
+        )
         self._compare_models(
-            m1, m2, is_deterministic=is_deterministic, use_virtual_table=True
+            m1,
+            m2,
+            is_deterministic=is_deterministic,
+            use_virtual_table=True,
+            tolerance=1e-2,
         )
 
     @unittest.skipIf(