Fix empty sharding constraints in test_model_parallel.py (#2998)

Summary:
Pull Request resolved: #2998

#### Context

Several unit tests in `test_model_parallel.py` passed **empty constraints** into `self._test_sharding` because the constraints are generated using an empty `self.tables` before invoking `self._build_tables_and_groups`.

Impacted tests are:
* `test_sharding_twcw`
* `test_sharding_variable_batch`
* `test_sharding_multiple_kernels`

#### Changes

* Constraints only depend on table names. A new list `self.table_names` is created in `setUp()` stage to be used to construct constraints.
* Updates `self._build_tables_and_groups` to use the generated table names.
* Increases `max_examples` for `test_sharding_multiple_kernels` to cover both FP32 and FP16 cases.

Reviewed By: TroyGarden

Differential Revision: D75306149

fbshipit-source-id: b93f7656e45a8c79393a1c347437f757aac07557

Author: Shuangping Liu

torchrec/distributed/test_utils/test_model_parallel.py

@@ -40,6 +40,10 @@ def setUp(self, backend: str = "nccl") -> None:
         self.num_weighted_features = 2
         self.num_shared_features = 2
 
+        self.table_names = [
+            "table_" + str(i)
+            for i in range(self.num_features + self.num_shared_features)
+        ]
         self.tables = []
         self.mean_tables = []
         self.weighted_tables = []
@@ -63,7 +67,7 @@ def _build_tables_and_groups(
             EmbeddingBagConfig(
                 num_embeddings=(i + 1) * 10,
                 embedding_dim=(i + 2) * 8,
-                name="table_" + str(i),
+                name=self.table_names[i],
                 feature_names=["feature_" + str(i)],
                 data_type=data_type,
             )
@@ -73,7 +77,7 @@ def _build_tables_and_groups(
             EmbeddingBagConfig(
                 num_embeddings=(i + 1) * 10,
                 embedding_dim=(i + 2) * 8,
-                name="table_" + str(i + self.num_features),
+                name=self.table_names[i + self.num_features],
                 feature_names=["feature_" + str(i)],
                 data_type=data_type,
             )
@@ -85,7 +89,7 @@ def _build_tables_and_groups(
             EmbeddingBagConfig(
                 num_embeddings=(i + 1) * 10,
                 embedding_dim=(i + 2) * 8,
-                name="table_" + str(i),
+                name=self.table_names[i],
                 feature_names=["feature_" + str(i)],
                 pooling=PoolingType.MEAN,
                 data_type=data_type,
@@ -97,7 +101,7 @@ def _build_tables_and_groups(
             EmbeddingBagConfig(
                 num_embeddings=(i + 1) * 10,
                 embedding_dim=(i + 2) * 8,
-                name="table_" + str(i + self.num_features),
+                name=self.table_names[i + self.num_features],
                 feature_names=["feature_" + str(i)],
                 pooling=PoolingType.MEAN,
                 data_type=data_type,
@@ -385,8 +389,8 @@ def test_sharding_cw(
             backend=self.backend,
             qcomms_config=qcomms_config,
             constraints={
-                table.name: ParameterConstraints(min_partition=4)
-                for table in self.tables
+                table_name: ParameterConstraints(min_partition=4)
+                for table_name in self.table_names
             },
             apply_optimizer_in_backward_config=apply_optimizer_in_backward_config,
             variable_batch_size=variable_batch_size,
@@ -466,8 +470,8 @@ def test_sharding_twcw(
             backend=self.backend,
             qcomms_config=qcomms_config,
             constraints={
-                table.name: ParameterConstraints(min_partition=4)
-                for table in self.tables
+                table_name: ParameterConstraints(min_partition=4)
+                for table_name in self.table_names
             },
             apply_optimizer_in_backward_config=apply_optimizer_in_backward_config,
             variable_batch_size=variable_batch_size,
@@ -681,8 +685,8 @@ def test_sharding_variable_batch(
             ],
             backend=self.backend,
             constraints={
-                table.name: ParameterConstraints(min_partition=4)
-                for table in self.tables
+                table_name: ParameterConstraints(min_partition=4)
+                for table_name in self.table_names
             },
             variable_batch_per_feature=True,
             has_weighted_tables=False,
@@ -700,24 +704,25 @@ def test_sharding_variable_batch(
         sharding_type=st.just(ShardingType.COLUMN_WISE.value),
         data_type=st.sampled_from([DataType.FP32, DataType.FP16]),
     )
-    @settings(verbosity=Verbosity.verbose, max_examples=1, deadline=None)
+    @settings(verbosity=Verbosity.verbose, max_examples=2, deadline=None)
     def test_sharding_multiple_kernels(
         self, sharding_type: str, data_type: DataType
     ) -> None:
         if self.backend == "gloo":
             self.skipTest("ProcessGroupGloo does not support reduce_scatter")
+        fused_params = {"prefetch_pipeline": True}
         constraints = {
-            table.name: ParameterConstraints(
+            table_name: ParameterConstraints(
                 min_partition=4,
                 compute_kernels=(
                     [EmbeddingComputeKernel.FUSED.value]
                     if i % 2 == 0
                     else [EmbeddingComputeKernel.FUSED_UVM_CACHING.value]
                 ),
+                sharding_types=[sharding_type],
             )
-            for i, table in enumerate(self.tables)
+            for i, table_name in enumerate(self.table_names)
         }
-        fused_params = {"prefetch_pipeline": True}
         self._test_sharding(
             # pyre-ignore[6]
             sharders=[EmbeddingBagCollectionSharder(fused_params=fused_params)],