fix test

Differential Revision: D81996413

Author: jd86-aai

torchrec/distributed/tests/test_init_parameters.py

@@ -87,6 +87,9 @@ def initialize_and_test_parameters(
             else f"embedding_bags.{table_name}.weight"
         )
 
+        # Store the original tensor on CPU for comparison
+        original_tensor = embedding_tables.state_dict()[key].clone().cpu()
+
         if isinstance(model.state_dict()[key], DTensor):
             if ctx.rank == 0:
                 gathered_tensor = torch.empty(model.state_dict()[key].size())
@@ -96,28 +99,26 @@ def initialize_and_test_parameters(
             gathered_tensor = model.state_dict()[key].full_tensor()
             if ctx.rank == 0:
                 torch.testing.assert_close(
-                    gathered_tensor,
-                    embedding_tables.state_dict()[key],
+                    gathered_tensor.cpu(), original_tensor, rtol=1e-5, atol=1e-6
                 )
         elif isinstance(model.state_dict()[key], ShardedTensor):
             if ctx.rank == 0:
-                gathered_tensor = torch.empty_like(
-                    embedding_tables.state_dict()[key], device=ctx.device
-                )
+                gathered_tensor = torch.empty_like(original_tensor, device=ctx.device)
             else:
                 gathered_tensor = None
 
             model.state_dict()[key].gather(dst=0, out=gathered_tensor)
 
             if ctx.rank == 0:
                 torch.testing.assert_close(
-                    none_throws(gathered_tensor).to("cpu"),
-                    embedding_tables.state_dict()[key],
+                    none_throws(gathered_tensor).cpu(),
+                    original_tensor,
+                    rtol=1e-5,
+                    atol=1e-6,
                 )
         elif isinstance(model.state_dict()[key], torch.Tensor):
             torch.testing.assert_close(
-                embedding_tables.state_dict()[key].cpu(),
-                model.state_dict()[key].cpu(),
+                model.state_dict()[key].cpu(), original_tensor, rtol=1e-5, atol=1e-6
             )
         else:
             raise AssertionError(
@@ -161,6 +162,9 @@ def test_initialize_parameters_ec(self, sharding_type: str) -> None:
         backend = "nccl"
         table_name = "free_parameters"
 
+        # Set seed for deterministic tensor generation
+        torch.manual_seed(42)
+
         # Initialize embedding table on non-meta device, in this case cuda:0
         embedding_tables = EmbeddingCollection(
             tables=[
@@ -173,8 +177,10 @@ def test_initialize_parameters_ec(self, sharding_type: str) -> None:
             ],
         )
 
+        # Use a fixed tensor with explicit seeding for consistent testing
+        fixed_tensor = torch.randn(10, 64, generator=torch.Generator().manual_seed(42))
         embedding_tables.load_state_dict(
-            {f"embeddings.{table_name}.weight": torch.randn(10, 64)}
+            {f"embeddings.{table_name}.weight": fixed_tensor}
         )
 
         self._run_multi_process_test(
@@ -210,6 +216,9 @@ def test_initialize_parameters_ebc(self, sharding_type: str) -> None:
         backend = "nccl"
         table_name = "free_parameters"
 
+        # Set seed for deterministic tensor generation
+        torch.manual_seed(42)
+
         # Initialize embedding bag on non-meta device, in this case cuda:0
         embedding_tables = EmbeddingBagCollection(
             tables=[
@@ -222,8 +231,10 @@ def test_initialize_parameters_ebc(self, sharding_type: str) -> None:
             ],
         )
 
+        # Use a fixed tensor with explicit seeding for consistent testing
+        fixed_tensor = torch.randn(10, 64, generator=torch.Generator().manual_seed(42))
         embedding_tables.load_state_dict(
-            {f"embedding_bags.{table_name}.weight": torch.randn(10, 64)}
+            {f"embedding_bags.{table_name}.weight": fixed_tensor}
         )
 
         self._run_multi_process_test(