From 7dc9d445b2e047f707d1ea6d0ed71f7bc4fcc6ab Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Fri, 15 Apr 2022 01:04:17 +0000
Subject: [PATCH] added basic distributed logic for trainer

---
 mattport/nerf/trainer.py                 | 58 ++++++++++++++++++++++--
 scripts/run_train_nerf.py                |  6 +--
 tests/nerf/configs/test_basic_graph.yaml |  1 +
 tests/nerf/test_trainer.py               | 15 ++++++
 4 files changed, 71 insertions(+), 9 deletions(-)
 create mode 100644 tests/nerf/test_trainer.py

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index b827960a4..4bf367231 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,6 +1,16 @@
 """
 Code to train model.
 """
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from omegaconf import DictConfig
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data.distributed import DistributedSampler
+
 from mattport.nerf.dataset.build import build_dataset
 from mattport.nerf.optimizer import Optimizer
 
@@ -44,14 +54,52 @@ class Trainer:
         """
         raise NotImplementedError
 
-    def train(self):
+    def train(self, gpu: int) -> None:
+        """_summary_"""
+        self.setup_ddp(gpu, self.config.world_size)
+
+        # TODO(): implement training pipeline. below I add skeleton code for compatibility with DDP
+        # self.setup_graph()
+        # self.graph = DDP(self.graph, device_ids=[gpu])
+        # self.setup_dataset()
+        # train_sampler = DistributedSampler(self.dataset, rank=gpu, num_replicas=self.config.world_size, shuffle=True)
+        # TODO(): create dataloader that takes as input the train_sampler
+        # self.setup_optimizer()
+
+        # TODO(): calls to train_epoch and train_iteration loop
+
+        dist.destroy_process_group()
+        raise NotImplementedError
+
+    def train_epoch(self, gpu: int):
         """_summary_"""
         raise NotImplementedError
 
-    def train_epoch(self):
+    def train_iteration(self, gpu: int):
         """_summary_"""
+        # TODO(): save checkpoints, and do logging only on rank 0 device only
+        # if gpu == 0:
         raise NotImplementedError
 
-    def train_iteration(self):
-        """_summary_"""
-        raise NotImplementedError
+    def setup_ddp(self, gpu: int, world_size: int) -> None:
+        """Initialization of processes for distributed training.
+
+        Args:
+            gpu (int): process rank number
+            world_size (int): total number of gpus available for training
+        """
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12356"
+        dist.init_process_group("nccl", init_method="env://", rank=gpu, world_size=world_size)
+        torch.cuda.set_device(gpu)
+
+        torch.manual_seed(0)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        np.random.seed(0)
+
+    def run_multiprocess_train(self) -> None:
+        """function to spawn multiple processes for training."""
+        if self.config.world_size < 1:
+            self.config.world_size = torch.cuda.device_count()
+        mp.spawn(self.train, nprocs=self.config.world_size)
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 2ad9a5f97..791020ac1 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -2,12 +2,12 @@
 run_train_nerf.py
 """
 import logging
+
 import hydra
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
 
-
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
 
@@ -15,9 +15,7 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin
 def main(config: DictConfig):
     """Main function."""
     trainer = Trainer(config)
-    trainer.setup_dataset()
-    trainer.setup_graph()
-    trainer.setup_optimizer()
+    trainer.run_multiprocess_train()
 
 
 if __name__ == "__main__":
diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/configs/test_basic_graph.yaml
index b3e23bd79..77939a274 100644
--- a/tests/nerf/configs/test_basic_graph.yaml
+++ b/tests/nerf/configs/test_basic_graph.yaml
@@ -1,3 +1,4 @@
+world_size: 0
 network:
   mlp_0:
     class_name: MLP
diff --git a/tests/nerf/test_trainer.py b/tests/nerf/test_trainer.py
new file mode 100644
index 000000000..0cc9ff4dd
--- /dev/null
+++ b/tests/nerf/test_trainer.py
@@ -0,0 +1,15 @@
+from tty import CFLAG
+import pytest
+import runpy
+
+from mattport.nerf.trainer import Trainer
+from hydra import compose, initialize
+
+CONFIG_DIR = "configs"
+
+def test_multiprocess_train():
+    """Test for trainer to validate distributed data parallel working."""
+    with initialize(config_path=CONFIG_DIR):
+        cfg = compose(config_name='test_basic_graph')
+    trainer = Trainer(cfg)
+    trainer.run_multiprocess_train()
\ No newline at end of file