add graph parallel initialization (#1032)

* add graph parallel initialization

* set default graph parallel group size to None

* set default graph parallel group size to None

* lint

* lint

src/fairchem/core/_cli_hydra.py

@@ -23,6 +23,8 @@
 from omegaconf import OmegaConf
 from omegaconf.errors import InterpolationKeyError
 
+from fairchem.core.common import gp_utils
+
 if TYPE_CHECKING:
     from omegaconf import DictConfig
 
@@ -109,6 +111,7 @@ class JobConfig:
     runner_state_path: Optional[str] = None  # noqa: UP007
     # read-only metadata about the job, not user inputs
     metadata: Optional[Metadata] = None  # noqa: UP007
+    graph_parallel_group_size: Optional[int] = None  # noqa: UP007 python 3.9 requires Optional still
 
     def __post_init__(self) -> None:
         self.metadata = Metadata(
@@ -154,7 +157,15 @@ def __call__(self, dict_config: DictConfig) -> None:
         # TODO also load job config here
         setup_env_vars()
         setup_logging()
-        distutils.setup(map_job_config_to_dist_config(self.config.job))
+
+        dist_config = map_job_config_to_dist_config(self.config.job)
+        distutils.setup(dist_config)
+        if self.config.job.graph_parallel_group_size is not None:
+            gp_utils.setup_graph_parallel_groups(
+                self.config.job.graph_parallel_group_size,
+                dist_config["distributed_backend"],
+            )
+
         self._init_logger()
         _set_seeds(self.config.job.seed)
         if self.config.job.deterministic:

src/fairchem/core/common/gp_utils.py

@@ -35,6 +35,42 @@ def divide_and_check_no_remainder(a: int, b: int) -> int:
     return a // b
 
 
+def setup_graph_parallel_groups(
+    graph_parallel_group_size: int, distributed_backend: str
+) -> None:
+    assert torch.distributed.is_initialized()
+    world_size = torch.distributed.get_world_size()
+    assert (
+        graph_parallel_group_size <= world_size
+    ), "graph parallel group size must be at most world size"
+
+    ensure_div(world_size, graph_parallel_group_size)
+    dp_size = world_size // graph_parallel_group_size
+    rank = dist.get_rank()
+
+    if rank == 0:
+        logging.info(
+            f"> initializing graph parallel with size {graph_parallel_group_size}"
+        )
+        logging.info(f"> initializing ddp with size {dp_size}")
+
+    groups = torch.arange(world_size).reshape(dp_size, graph_parallel_group_size)
+    found = [x.item() for x in torch.where(groups == rank)]
+
+    global _DATA_PARALLEL_GROUP
+    assert _DATA_PARALLEL_GROUP is None, "data parallel group is already initialized"
+    for j in range(graph_parallel_group_size):
+        group = dist.new_group(groups[:, j].tolist(), backend=distributed_backend)
+        if j == found[1]:
+            _DATA_PARALLEL_GROUP = group
+    global _GRAPH_PARALLEL_GROUP
+    assert _GRAPH_PARALLEL_GROUP is None, "graph parallel group is already initialized"
+    for i in range(dp_size):
+        group = dist.new_group(groups[i, :].tolist(), backend=distributed_backend)
+        if i == found[0]:
+            _GRAPH_PARALLEL_GROUP = group
+
+
 def setup_gp(config) -> None:
     gp_size = config["gp_gpus"]
     backend = config["distributed_backend"]