Add barrier around dataset processor for race condition

Author: szaman19

DGraph/distributed/nccl/NCCLBackendEngine.py

@@ -512,6 +512,11 @@ def __init__(self, ranks_per_graph=-1, *args, **kwargs):
         if not NCCLBackendEngine._is_initialized:
             self.init_process_group(ranks_per_graph)
 
+    def barrier(self) -> None:
+        if not dist.is_initialized():
+            raise RuntimeError("NCCL backend engine is not initialized")
+        dist.barrier()
+
     def init_process_group(self, ranks_per_graph=-1, *args, **kwargs):
         if not dist.is_initialized():
             dist.init_process_group(backend="nccl", *args, **kwargs)

experiments/OGB/GCN.py

@@ -14,6 +14,7 @@
 import torch
 import torch.nn as nn
 import torch.distributed as dist
+from DGraph.utils.TimingReport import TimingReport
 
 
 class ConvLayer(nn.Module):
@@ -54,24 +55,41 @@ def forward(
         num_local_nodes = node_features.size(1)
         _src_indices = edge_index[:, 0, :]
         _dst_indices = edge_index[:, 1, :]
+        TimingReport.start("pre-processing")
         _src_rank_mappings = torch.cat(
             [rank_mapping[0].unsqueeze(0), rank_mapping[0].unsqueeze(0)], dim=0
         )
         _dst_rank_mappings = torch.cat(
             [rank_mapping[0].unsqueeze(0), rank_mapping[1].unsqueeze(0)], dim=0
         )
+        TimingReport.stop("pre-processing")
+        TimingReport.start("Gather_1")
         x = self.comm.gather(
             node_features, _dst_indices, _dst_rank_mappings, cache=gather_cache
         )
+        TimingReport.stop("Gather_1")
+        TimingReport.start("Conv_1")
         x = self.conv1(x)
+        TimingReport.stop("Conv_1")
+        TimingReport.start("Scatter_1")
         x = self.comm.scatter(
             x, _src_indices, _src_rank_mappings, num_local_nodes, cache=scatter_cache
         )
+        TimingReport.stop("Scatter_1")
+        TimingReport.start("Gather_2")
         x = self.comm.gather(x, _dst_indices, _dst_rank_mappings, cache=gather_cache)
+        TimingReport.stop("Gather_2")
+        TimingReport.start("Conv_2")
         x = self.conv2(x)
+        TimingReport.stop("Conv_2")
+        TimingReport.start("Scatter_2")
         x = self.comm.scatter(
             x, _src_indices, _src_rank_mappings, num_local_nodes, cache=scatter_cache
         )
+        TimingReport.stop("Scatter_2")
+        TimingReport.start("Final_FC")
         x = self.fc(x)
+        TimingReport.stop("Final_FC")
+
         # x = self.softmax(x)
         return x

experiments/OGB/main.py

@@ -38,6 +38,8 @@
 )
 import numpy as np
 import os
+from DGraph.utils.TimingReport import TimingReport
+import json
 
 
 class SingleProcessDummyCommunicator(CommunicatorBase):
@@ -131,7 +133,6 @@ def _run_experiment(
             print(f"Rank: {rank} Mapping: {rank_mappings.shape}")
             print(f"Rank: {rank} Node Features: {node_features.shape}")
             print(f"Rank: {rank} Edge Indices: {edge_indices.shape}")
-
         comm.barrier()
     criterion = torch.nn.CrossEntropyLoss()
 
@@ -229,7 +230,9 @@ def _run_experiment(
             assert rank != rank
             assert value.shape[0] == scatter_cache.gather_recv_comm_vector
         end_time = perf_counter()
-        print(f"Rank: {rank} Cache Generation Time: {end_time - start_time:.4f} s")
+        elapsed_time_in_ms = (end_time - start_time) * 1000
+        print(f"Rank: {rank} Cache Generation Time: {elapsed_time_in_ms:.4f} ms")
+        TimingReport.add_time("cache_generation_time", elapsed_time_in_ms)
 
         # with open(f"{log_prefix}_gather_cache_{world_size}_{rank}.pt", "wb") as f:
         #    torch.save(gather_cache, f)
@@ -366,6 +369,7 @@ def main(
                 node_rank_placement_file, weights_only=False
             )
 
+    TimingReport.init(comm)
     safe_create_dir(log_dir, comm.get_rank())
     training_dataset = DistributedOGBWrapper(
         f"ogbn-{dataset}",
@@ -381,7 +385,7 @@ def main(
     validation_accuracies = np.zeros((runs, epochs))
     world_size = comm.get_world_size()
 
-    dist.barrier()
+    comm.barrier()
     print(f"Running experiment with {world_size} processes on dataset {dataset}")
     print(f"Using cache: {use_cache}")
 
@@ -402,6 +406,11 @@ def main(
         validation_trajectores[i] = val_traj
         validation_accuracies[i] = val_accuracy
 
+    write_experiment_log(
+        json.dumps(TimingReport._timers),
+        f"{log_dir}/timing_report_world_size_{world_size}_cache_{use_cache}.json",
+        comm.get_rank(),
+    )
     visualize_trajectories(
         training_trajectores,
         "Training Loss",