Add numerical equivalence tests for model_parallel (#1721)

Summary:
Pull Request resolved: #1721

Adding numerical equivalence tests for different kernel types.

Note that when stochastic rounding = True + is_training = True is always a bit tricky. However, from what FBGEMM team said, it is not guaranteed to have iteration-wise equivalence, but it's guaranteed to have numerical equivalence in the long term.

Reviewed By: sarckk

Differential Revision: D53878573

fbshipit-source-id: 5b0eb4f1977a448f27db80e9bd19b02cb58b15bb

Author: henrylhtsang

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -52,6 +52,7 @@
 )
 from torchrec.modules.embedding_configs import (
     BaseEmbeddingConfig,
+    DataType,
     EmbeddingBagConfig,
     PoolingType,
 )
@@ -240,7 +241,7 @@ def test_sharding_fused_ebc_as_top_level(self) -> None:
         self.assertTrue(isinstance(model.module, ShardedFusedEmbeddingBagCollection))
 
 
-class ModelParallelStateDictBase(unittest.TestCase):
+class ModelParallelSingleRankBase(unittest.TestCase):
     def setUp(self, backend: str = "nccl") -> None:
         os.environ["RANK"] = "0"
         os.environ["WORLD_SIZE"] = "1"
@@ -261,9 +262,79 @@ def setUp(self, backend: str = "nccl") -> None:
 
         dist.init_process_group(backend=backend)
 
+    def tearDown(self) -> None:
+        dist.destroy_process_group()
+        del os.environ["NCCL_SOCKET_IFNAME"]
+        super().tearDown()
+
+    def _train_models(
+        self,
+        m1: DistributedModelParallel,
+        m2: DistributedModelParallel,
+        batch: ModelInput,
+    ) -> None:
+        loss1, pred1 = m1(batch)
+        loss2, pred2 = m2(batch)
+        loss1.backward()
+        loss2.backward()
+
+    def _eval_models(
+        self,
+        m1: DistributedModelParallel,
+        m2: DistributedModelParallel,
+        batch: ModelInput,
+        is_deterministic: bool = True,
+    ) -> 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)
+
+    def _compare_models(
+        self,
+        m1: DistributedModelParallel,
+        m2: DistributedModelParallel,
+        is_deterministic: bool = True,
+    ) -> None:
+        sd1 = m1.state_dict()
+        for key, value in m2.state_dict().items():
+            v2 = sd1[key]
+            if isinstance(value, ShardedTensor):
+                assert isinstance(v2, ShardedTensor)
+                self.assertEqual(len(value.local_shards()), len(v2.local_shards()))
+                for dst, src in zip(value.local_shards(), v2.local_shards()):
+                    if is_deterministic:
+                        self.assertTrue(torch.equal(src.tensor, dst.tensor))
+                    else:
+                        rtol, atol = _get_default_rtol_and_atol(src.tensor, dst.tensor)
+                        torch.testing.assert_close(
+                            src.tensor, dst.tensor, rtol=rtol, atol=atol
+                        )
+            else:
+                dst = value
+                src = v2
+                if is_deterministic:
+                    self.assertTrue(torch.equal(src, dst))
+                else:
+                    rtol, atol = _get_default_rtol_and_atol(src, dst)
+                    torch.testing.assert_close(src, dst, rtol=rtol, atol=atol)
+
+
+class ModelParallelStateDictBase(ModelParallelSingleRankBase):
+    def setUp(self, backend: str = "nccl") -> None:
+        super().setUp(backend=backend)
+
         num_features = 4
         num_weighted_features = 2
-        self.batch_size = 3
+        self.batch_size = 20
         self.num_float_features = 10
 
         self.tables = [
@@ -285,11 +356,6 @@ def setUp(self, backend: str = "nccl") -> None:
             for i in range(num_weighted_features)
         ]
 
-    def tearDown(self) -> None:
-        dist.destroy_process_group()
-        del os.environ["NCCL_SOCKET_IFNAME"]
-        super().tearDown()
-
     def _generate_dmps_and_batch(
         self,
         sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
@@ -352,6 +418,13 @@ def _generate_dmps_and_batch(
             dmps.append(dmp)
         return (dmps, batch)
 
+    def _set_table_weights_precision(self, dtype: DataType) -> None:
+        for table in self.tables:
+            table.data_type = dtype
+
+        for weighted_table in self.weighted_tables:
+            weighted_table.data_type = dtype
+
     def test_parameter_init(self) -> None:
         class MyModel(nn.Module):
             def __init__(self, device: str, val: float) -> None:
@@ -513,33 +586,9 @@ def test_load_state_dict(
 
         # validate the models are equivalent
         if is_training:
-            for _ in range(2):
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                loss1.backward()
-                loss2.backward()
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        else:
-            with torch.no_grad():
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        sd1 = m1.state_dict()
-        for key, value in m2.state_dict().items():
-            v2 = sd1[key]
-            if isinstance(value, ShardedTensor):
-                assert len(value.local_shards()) == 1
-                dst = value.local_shards()[0].tensor
-            else:
-                dst = value
-            if isinstance(v2, ShardedTensor):
-                assert len(v2.local_shards()) == 1
-                src = v2.local_shards()[0].tensor
-            else:
-                src = v2
-            self.assertTrue(torch.equal(src, dst))
+            self._train_models(m1, m2, batch)
+        self._eval_models(m1, m2, batch)
+        self._compare_models(m1, m2)
 
     # pyre-ignore[56]
     @given(
@@ -582,33 +631,9 @@ def test_load_state_dict_dp(
 
         # validate the models are equivalent
         if is_training:
-            for _ in range(2):
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                loss1.backward()
-                loss2.backward()
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        else:
-            with torch.no_grad():
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        sd1 = m1.state_dict()
-        for key, value in m2.state_dict().items():
-            v2 = sd1[key]
-            if isinstance(value, ShardedTensor):
-                assert len(value.local_shards()) == 1
-                dst = value.local_shards()[0].tensor
-            else:
-                dst = value
-            if isinstance(v2, ShardedTensor):
-                assert len(v2.local_shards()) == 1
-                src = v2.local_shards()[0].tensor
-            else:
-                src = v2
-            self.assertTrue(torch.equal(src, dst))
+            self._train_models(m1, m2, batch)
+        self._eval_models(m1, m2, batch)
+        self._compare_models(m1, m2)
 
     # pyre-ignore[56]
     @given(
@@ -661,34 +686,9 @@ def test_load_state_dict_prefix(
 
         # validate the models are equivalent
         if is_training:
-            for _ in range(2):
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                loss1.backward()
-                loss2.backward()
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        else:
-            with torch.no_grad():
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-
-        sd1 = m1.state_dict()
-        for key, value in m2.state_dict().items():
-            v2 = sd1[key]
-            if isinstance(value, ShardedTensor):
-                assert len(value.local_shards()) == 1
-                dst = value.local_shards()[0].tensor
-            else:
-                dst = value
-            if isinstance(v2, ShardedTensor):
-                assert len(v2.local_shards()) == 1
-                src = v2.local_shards()[0].tensor
-            else:
-                src = v2
-            self.assertTrue(torch.equal(src, dst))
+            self._train_models(m1, m2, batch)
+        self._eval_models(m1, m2, batch)
+        self._compare_models(m1, m2)
 
     # pyre-fixme[56]
     @given(
@@ -807,19 +807,8 @@ def test_load_state_dict_cw_multiple_shards(
 
         # validate the models are equivalent
         if is_training:
-            for _ in range(2):
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                loss1.backward()
-                loss2.backward()
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
-        else:
-            with torch.no_grad():
-                loss1, pred1 = m1(batch)
-                loss2, pred2 = m2(batch)
-                self.assertTrue(torch.equal(loss1, loss2))
-                self.assertTrue(torch.equal(pred1, pred2))
+            self._train_models(m1, m2, batch)
+        self._eval_models(m1, m2, batch)
 
         sd1 = m1.state_dict()
         for key, value in m2.state_dict().items():
@@ -876,3 +865,89 @@ def test_load_state_dict_cw_multiple_shards(
                         )
                 elif isinstance(dst_opt_state, torch.Tensor):
                     self.assertIsInstance(src_opt_state, torch.Tensor)
+
+    @unittest.skipIf(
+        not torch.cuda.is_available(),
+        "Not enough GPUs, this test requires at least one GPU",
+    )
+    # pyre-ignore[56]
+    @given(
+        sharder_type=st.sampled_from(
+            [
+                SharderType.EMBEDDING_BAG_COLLECTION.value,
+            ]
+        ),
+        sharding_type=st.sampled_from(
+            [
+                ShardingType.TABLE_WISE.value,
+                ShardingType.COLUMN_WISE.value,
+                ShardingType.ROW_WISE.value,
+                ShardingType.TABLE_ROW_WISE.value,
+                ShardingType.TABLE_COLUMN_WISE.value,
+            ]
+        ),
+        kernel_type=st.sampled_from(
+            [
+                EmbeddingComputeKernel.FUSED_UVM_CACHING.value,
+                EmbeddingComputeKernel.FUSED_UVM.value,
+            ]
+        ),
+        is_training=st.booleans(),
+        stochastic_rounding=st.booleans(),
+        dtype=st.sampled_from([DataType.FP32, DataType.FP16]),
+    )
+    @settings(verbosity=Verbosity.verbose, max_examples=4, deadline=None)
+    def test_numerical_equivalence_between_kernel_types(
+        self,
+        sharder_type: str,
+        sharding_type: str,
+        kernel_type: str,
+        is_training: bool,
+        stochastic_rounding: bool,
+        dtype: DataType,
+    ) -> None:
+        self._set_table_weights_precision(dtype)
+        fused_params = {
+            "stochastic_rounding": stochastic_rounding,
+            "cache_precision": dtype,
+        }
+
+        fused_sharders = [
+            cast(
+                ModuleSharder[nn.Module],
+                create_test_sharder(
+                    sharder_type,
+                    sharding_type,
+                    EmbeddingComputeKernel.FUSED.value,
+                    fused_params=fused_params,
+                ),
+            ),
+        ]
+        sharders = [
+            cast(
+                ModuleSharder[nn.Module],
+                create_test_sharder(
+                    sharder_type,
+                    sharding_type,
+                    kernel_type,
+                    fused_params=fused_params,
+                ),
+            ),
+        ]
+        (fused_model, _), _ = self._generate_dmps_and_batch(fused_sharders)
+        (model, _), batch = self._generate_dmps_and_batch(sharders)
+
+        # load the baseline model's state_dict onto the new model
+        model.load_state_dict(
+            cast("OrderedDict[str, torch.Tensor]", fused_model.state_dict())
+        )
+
+        if is_training:
+            for _ in range(4):
+                self._train_models(fused_model, model, batch)
+        self._eval_models(
+            fused_model, model, batch, is_deterministic=not stochastic_rounding
+        )
+        self._compare_models(
+            fused_model, model, is_deterministic=not stochastic_rounding
+        )