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

Summary:

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

Differential Revision: D77742701

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,11 @@ 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)
         dedup_factor = 2
 
         global_kjt, local_kjts = ModelInput._generate_variable_batch_features(

torchrec/distributed/test_utils/test_sharding.py

@@ -135,6 +135,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"]]: ...
 
 
@@ -152,6 +153,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,8 +182,10 @@ 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)
     if dedup_feature_names:
         model = model_class(
             tables=cast(
@@ -224,6 +228,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:
@@ -242,6 +247,7 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
+                    random_seed=random_seed,
                 )
             )
     else:
@@ -259,6 +265,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)
@@ -718,6 +725,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: int = 0,
 ) -> None:
     batch_size = random.randint(0, batch_size) if allow_zero_batch_size else batch_size
     # Generate model & inputs.
@@ -746,7 +754,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)
@@ -794,6 +804,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.
@@ -973,6 +984,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: int = 100,
 ) -> None:
     with MultiProcessContext(rank, world_size, backend, local_size) as ctx:
         assert ctx.pg is not None
@@ -1006,6 +1018,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

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