Add random_seed for regular model parallel tests to ensure actual randomness in generating embeddings/inputs etc... (#3158)

Summary:
Pull Request resolved: #3158

Add random_seed as an optional parameter for gen_model_and_input method that can be used by any other testing methods.

Reviewed By: aporialiao

Differential Revision: D77742701

fbshipit-source-id: 752c7a9fd84436a5de862413f1851cf09a25d38e

Author: isururanawaka

torchrec/distributed/test_utils/test_model.py

@@ -100,11 +100,14 @@ def generate(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
+        random_seed: Optional[int] = None,
     ) -> Tuple["ModelInput", List["ModelInput"]]:
         """
         Returns a global (single-rank training) batch
         and a list of local (multi-rank training) batches of world_size.
         """
+        if random_seed is not None:
+            torch.manual_seed(random_seed)
         batch_size_by_rank = [batch_size] * world_size
         if variable_batch_size:
             batch_size_by_rank = [
@@ -751,9 +754,14 @@ def generate_variable_batch_input(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
+        random_seed: Optional[int] = None,
     ) -> Tuple["ModelInput", List["ModelInput"]]:
-        torch.manual_seed(100)
-        random.seed(100)
+        if random_seed is not None:
+            torch.manual_seed(random_seed)
+            random.seed(random_seed)
+        else:
+            torch.manual_seed(100)
+            random.seed(100)
         dedup_factor = 2
 
         global_kjt, local_kjts = ModelInput._generate_variable_batch_features(

torchrec/distributed/test_utils/test_sharding.py

@@ -163,6 +163,7 @@ def __call__(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
+        random_seed: Optional[int] = None,
     ) -> Tuple["ModelInput", List["ModelInput"]]: ...
 
 
@@ -180,6 +181,7 @@ def __call__(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
+        random_seed: Optional[int] = None,
     ) -> Tuple["ModelInput", List["ModelInput"]]: ...
 
 
@@ -208,8 +210,12 @@ def gen_model_and_input(
     global_constant_batch: bool = False,
     num_inputs: int = 1,
     input_type: str = "kjt",  # "kjt" or "td"
+    random_seed: Optional[int] = None,
 ) -> Tuple[nn.Module, List[Tuple[ModelInput, List[ModelInput]]]]:
-    torch.manual_seed(0)
+    if random_seed is not None:
+        torch.manual_seed(random_seed)
+    else:
+        torch.manual_seed(0)
     if dedup_feature_names:
         model = model_class(
             tables=cast(
@@ -252,6 +258,7 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
+                    random_seed=random_seed,
                 )
             )
     elif generate == ModelInput.generate:
@@ -270,6 +277,7 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
+                    random_seed=random_seed,
                 )
             )
     else:
@@ -287,6 +295,7 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
+                    random_seed=random_seed,
                 )
             )
     return (model, inputs)
@@ -742,6 +751,7 @@ def sharding_single_rank_test_single_process(
     indices_dtype: torch.dtype = torch.int64,
     offsets_dtype: torch.dtype = torch.int64,
     lengths_dtype: torch.dtype = torch.int64,
+    random_seed: Optional[int] = None,
 ) -> None:
     batch_size = random.randint(0, batch_size) if allow_zero_batch_size else batch_size
     # Generate model & inputs.
@@ -770,7 +780,9 @@ def sharding_single_rank_test_single_process(
         indices_dtype=indices_dtype,
         offsets_dtype=offsets_dtype,
         lengths_dtype=lengths_dtype,
+        random_seed=random_seed,
     )
+
     global_model = global_model.to(device)
     global_input = inputs[0][0].to(device)
     local_input = inputs[0][1][rank].to(device)
@@ -818,6 +830,7 @@ def sharding_single_rank_test_single_process(
         constraints=constraints,
     )
     plan: ShardingPlan = planner.collective_plan(local_model, sharders, pg)
+
     """
     Simulating multiple nodes on a single node. However, metadata information and
     tensor placement must still be consistent. Here we overwrite this to do so.
@@ -994,6 +1007,7 @@ def sharding_single_rank_test(
     indices_dtype: torch.dtype = torch.int64,
     offsets_dtype: torch.dtype = torch.int64,
     lengths_dtype: torch.dtype = torch.int64,
+    random_seed: Optional[int] = None,
 ) -> None:
     with MultiProcessContext(rank, world_size, backend, local_size) as ctx:
         assert ctx.pg is not None
@@ -1027,6 +1041,7 @@ def sharding_single_rank_test(
             indices_dtype=indices_dtype,
             offsets_dtype=offsets_dtype,
             lengths_dtype=lengths_dtype,
+            random_seed=random_seed,
         )
 
 

torchrec/distributed/tests/test_model_parallel_nccl_ssd_single_gpu.py

@@ -609,6 +609,7 @@ def test_ssd_mixed_kernels_with_vbe(
             },
             constraints=constraints,
             variable_batch_per_feature=True,
+            random_seed=100,
         )
 
     @unittest.skipIf(