Merge pull request #22 from ethanweber/eng/distributed_training

Eng/distributed training

mattport/nerf/field_modules/mlp.py

@@ -39,7 +39,7 @@ class MLP(FieldModule):
         self.out_activation = out_activation
         self.net = None
 
-    def build_model(self) -> None:
+    def build_nn_modules(self) -> None:
         """Initialize mulilayer perceptron."""
         layers = []
         for i in range(self.num_layers):

mattport/nerf/graph.py

@@ -30,7 +30,7 @@ class Node:
         return hash(self.name)
 
 
-class Graph(nn.Module):
+class Graph(nn.ModuleDict):
     """_summary_"""
 
     def __init__(self, modules_config: DictConfig) -> None:
@@ -40,20 +40,19 @@ class Graph(nn.Module):
         self.roots = self.construct_graph()
         self.module_order = self.get_module_order()
 
-        self.modules = {}
         # initialize graph with known input dimensions; set default in_dim to 0
         for module_name, module_dict in modules_config.items():
             module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name)
             if not module_dict.meta_data.in_dim:
                 module_dict.meta_data.in_dim = 0
-            self.modules[module_name] = module(**module_dict.meta_data)
+            self[module_name] = module(**module_dict.meta_data)
 
         # calculate input dimensions based on module dependencies
         for root in self.roots:
             self.get_in_dim(root)
 
         # instantiate torch.nn members of network
-        for _, module in self.modules.items():
+        for _, module in self.items():
             module.build_nn_modules()
 
     def get_in_dim(self, curr_node: Node) -> None:
@@ -66,8 +65,8 @@ class Graph(nn.Module):
         if len(curr_node.parents) > 0:
             in_dim = 0
             for parent_name in curr_node.parents.keys():
-                in_dim += self.modules[parent_name].get_out_dim()
-            self.modules[curr_node.name].set_in_dim(in_dim)
+                in_dim += self[parent_name].get_out_dim()
+            self[curr_node.name].set_in_dim(in_dim)
             self.modules_config[curr_node.name].meta_data.in_dim = in_dim
 
         for child_node in curr_node.children.values():

mattport/nerf/trainer.py

@@ -1,13 +1,20 @@
 """
 Code to train model.
 """
+from omegaconf import DictConfig
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from mattport.nerf.graph import Graph
 from mattport.nerf.optimizer import Optimizer
 
 
 class Trainer:
     """Training class"""
 
-    def __init__(self, config: dict):
+    def __init__(self, local_rank: int, world_size: int, config: DictConfig):
+        self.local_rank = local_rank
+        self.world_size = world_size
         self.config = config
         self.train_dataset = None
         self.test_dataset = None
@@ -21,7 +28,10 @@ class Trainer:
 
     def setup_graph(self):
         """_summary_"""
-        raise NotImplementedError
+        self.graph = Graph(self.config.network).to(f"cuda:{self.local_rank}")
+        if self.world_size > 1:
+            self.graph = DDP(self.graph, device_ids=[self.local_rank])
+            dist.barrier(device_ids=[self.local_rank])
 
     def setup_optimizer(self):
         """_summary_"""
@@ -43,8 +53,9 @@ class Trainer:
         """
         raise NotImplementedError
 
-    def train(self):
+    def train(self) -> None:
         """_summary_"""
+        # is_main_thread = self.local_rank % self.world_size == 0
         raise NotImplementedError
 
     def train_epoch(self):
@@ -53,4 +64,6 @@ class Trainer:
 
     def train_iteration(self):
         """_summary_"""
+        # TODO(): save checkpoints, and do logging only on rank 0 device only
+        # if gpu == 0:
         raise NotImplementedError

scripts/run_train_nerf.py

@@ -1,23 +1,193 @@
 """
 run_train_nerf.py
 """
+import datetime
 import logging
+import random
+import socket
+from datetime import timedelta
+from typing import Any, Callable
+
 import hydra
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
 
-
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
+DEFAULT_TIMEOUT = timedelta(minutes=30)
+
+
+def _find_free_port() -> str:
+    """Find free port on open socket"""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def _set_random_seed(seed) -> None:
+    """Set randomness seed in torch and numpy"""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _synchronize(world_size, local_rank):
+    """Wait for all processes on node finish"""
+    if world_size > 1:
+        dist.barrier(device_ids=[local_rank])
+
+
+def _distributed_worker(
+    local_rank: int,
+    main_func: Callable,
+    world_size: int,
+    num_gpus_per_machine: int,
+    machine_rank: int,
+    dist_url: str,
+    config: DictConfig,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> Any:
+    """Spawned distributed worker that handles the initialization of process group and handles the
+       training process on multiple processes
+
+    Args:
+        local_rank (int): current rank of process
+        main_func (Callable): function that will be called by the distributed workers
+        world_size (int): total number of gpus available
+        num_gpus_per_machine (int): number of GPUs per machine
+        machine_rank (int): rank of this machine
+        dist_url (str): url to connect to for distributed jobs, including protocol
+                        e.g. "tcp://127.0.0.1:8686".
+                        Can be set to "auto" to automatically select a free port on localhost
+        config (DictConfig): config file specifying trainng regimen
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+
+    Raises:
+        e: Exception in initializing the process group
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    assert torch.cuda.is_available(), "cuda is not available. Please check your installation."
+    global_rank = machine_rank * num_gpus_per_machine + local_rank
+    try:
+        dist.init_process_group(
+            backend="NCCL",
+            init_method=dist_url,
+            world_size=world_size,
+            rank=global_rank,
+            timeout=timeout,
+        )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.error("Process group URL: %s", dist_url)
+        raise e
+
+    assert num_gpus_per_machine <= torch.cuda.device_count()
+    torch.cuda.set_device(local_rank)
+    _set_random_seed(1234 + local_rank)
+    _synchronize(world_size, local_rank)
+
+    output = main_func(local_rank, world_size, config)
+    _synchronize(world_size, local_rank)
+    dist.destroy_process_group()
+    return output
+
+
+def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
+    """Main training function that sets up and runs the trainer per process
+
+    Args:
+        local_rank (int): _description_
+        world_size (int): _description_
+        config (DictConfig): _description_
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    trainer = Trainer(local_rank, world_size, config)
+    trainer.setup_graph()
+    trainer.setup_optimizer()
+    trainer.setup_dataset()
+    trainer.train()
+    return 0
+
+
+def launch(
+    main_func: Callable,
+    num_gpus_per_machine: int,
+    num_machines: int = 1,
+    machine_rank: int = 0,
+    dist_url: str = "auto",
+    config: DictConfig = None,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> None:
+    """Function that spawns muliple processes to call on main_func
+
+    Args:
+        main_func (Callable): function that will be called by the distributed workers
+        num_gpus_per_machine (int): number of GPUs per machine
+        num_machines (int, optional): total number of machines
+        machine_rank (int, optional): rank of this machine. Defaults to 0.
+        dist_url (str, optional): url to connect to for distributed jobs. Defaults to "auto".
+        config (DictConfig, optional): config file specifying trainng regimen Defaults to None.
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+    """
+    world_size = num_machines * num_gpus_per_machine
+    if world_size > 1:
+        if dist_url == "auto":
+            assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs."
+            port = _find_free_port()
+            dist_url = f"tcp://127.0.0.1:{port}"
+        if num_machines > 1 and dist_url.startswith("file://"):
+            logger = logging.getLogger(__name__)
+            logger.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://")
+
+        process_context = mp.spawn(
+            _distributed_worker,
+            nprocs=num_gpus_per_machine,
+            args=(
+                main_func,
+                world_size,
+                num_gpus_per_machine,
+                machine_rank,
+                dist_url,
+                config,
+                timeout,
+            ),
+            join=False,
+        )
+        try:
+            process_context.join()
+        except KeyboardInterrupt:
+            for i, process in enumerate(process_context.processes):
+                if process.is_alive():
+                    logger.info("Terminating process %s", str(i))
+                    process.terminate()
+                process.join()
+                logger.info("Process %s finished", str(i))
+    else:
+        main_func(local_rank=0, world_size=1, config=config)
+
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
 def main(config: DictConfig):
     """Main function."""
-    trainer = Trainer(config)
-    trainer.setup_dataset()
-    trainer.setup_graph()
-    trainer.setup_optimizer()
+    launch(
+        _train,
+        config.machine_config.num_gpus,
+        num_machines=config.machine_config.num_machines,
+        machine_rank=config.machine_config.machine_rank,
+        dist_url=config.machine_config.dist_url,
+        config=config,
+    )
 
 
 if __name__ == "__main__":

tests/nerf/graph/configs/test_basic_graph.yaml

@@ -1,3 +1,8 @@
+machine_config:
+  num_gpus: 8
+  num_machines: 1
+  machine_rank: 0
+  dist_url: auto
 network:
   mlp_0:
     class_name: MLP

tests/nerf/graph/test_graph.py

@@ -32,17 +32,17 @@ def instantiate_graph(config_name: str) -> Graph:
 def test_graph_init():
     """test field modules dictionary is instantiated properly"""
     test_graph, _ = instantiate_graph("test_basic_graph")
-    assert test_graph.modules is not None
-    assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
-    assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
-    assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
+    assert test_graph is not None
+    assert type(test_graph["encoder_0"]).__name__ == "ScalingAndOffset"
+    assert type(test_graph["mlp_0"]).__name__ == "MLP"
+    assert type(test_graph["mlp_1"]).__name__ == "MLP"
 
 
 def test_input_dimension():
     """test field module input dimension calculation"""
     test_graph, _ = instantiate_graph("test_repeat_graph")
-    assert test_graph.modules["mlp_0"].in_dim == 16
-    assert test_graph.modules["mlp_1"].in_dim == 1
+    assert test_graph["mlp_0"].in_dim == 16
+    assert test_graph["mlp_1"].in_dim == 1
 
 
 def test_construct_graph_basic():