Minor fixes for review

Author: swong3-sc

python/tests/unit/module/models_test.py

@@ -299,13 +299,12 @@ def test_dmp_wrapped_model_produces_correct_output(self):
         """
         process_group_init_method = get_process_group_init_method()
         # Initialize distributed
-        if not dist.is_initialized():
-            dist.init_process_group(
-                backend="gloo",
-                init_method=process_group_init_method,
-                rank=0,
-                world_size=1,
-            )
+        dist.init_process_group(
+            backend="gloo",
+            init_method=process_group_init_method,
+            rank=0,
+            world_size=1,
+        )
 
         # Create model
         model = self._create_lightgcn_model(self.num_nodes)
@@ -336,13 +335,12 @@ def test_dmp_gradient_flow(self):
         process_group_init_method = get_process_group_init_method()
 
         # Initialize distributed
-        if not dist.is_initialized():
-            dist.init_process_group(
-                backend="gloo",
-                init_method=process_group_init_method,
-                rank=0,
-                world_size=1,
-            )
+        dist.init_process_group(
+            backend="gloo",
+            init_method=process_group_init_method,
+            rank=0,
+            world_size=1,
+        )
 
         # Create and wrap model
         model = self._create_lightgcn_model(self.num_nodes)
@@ -354,15 +352,13 @@ def test_dmp_gradient_flow(self):
 
         self._set_embeddings(model, "default_homogeneous_node_type")
 
-        model.train()
-
-        # Forward and backward pass
-        output = dmp_model(data=self.data, device=self.device)
+        dmp_model.train()
+        output = dmp_model(self.data, self.device)
         loss = output.sum()
         loss.backward()
 
         # Check that gradients exist and are non-zero
-        embedding_table = model._embedding_bag_collection.embedding_bags[
+        embedding_table = dmp_model._dmp_wrapped_module._embedding_bag_collection.embedding_bags[
             "node_embedding_default_homogeneous_node_type"
         ]
         self.assertIsNotNone(
@@ -383,22 +379,22 @@ def test_dmp_multiprocess(self, _name, world_size):
         """
         Test DMP with multiple processes to verify embedding sharding works correctly.
 
-        Note: Uses CPU/Gloo backend for unit testing. For GPU/NCCL testing, see integration tests.
+        Note: Uses CPU/Gloo backend for unit testing.
         """
         process_group_init_method = get_process_group_init_method()
 
         # Spawn world_size processes
         mp.spawn(
             fn=_run_dmp_multiprocess_test,
             args=(
-                world_size,
-                process_group_init_method,
-                self.num_nodes,
-                self.embedding_dim,
-                self.num_layers,
-                self.edge_index,
-                self.test_embeddings,
-                self.expected_output,
+                world_size, # total number of processes
+                process_group_init_method, # initialization method for process group
+                self.num_nodes, # number of nodes in test graph
+                self.embedding_dim, # dimension of embeddings
+                self.num_layers, # number of LightGCN layers
+                self.edge_index, # edge connectivity
+                self.test_embeddings, # test embedding values
+                self.expected_output, # expected model output
             ),
             nprocs=world_size,
         )
@@ -426,7 +422,7 @@ def _run_dmp_multiprocess_test(
         embedding_dim: Dimension of embeddings
         num_layers: Number of LightGCN layers
         edge_index: Edge connectivity
-        test_embeddings: Expected embedding values
+        test_embeddings: Test embedding values
         expected_output: Expected model output
     """
     try: