From ac3e7515ef416b719f27534b6c657bd62c833fcc Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 15 Apr 2022 21:33:50 +0000 Subject: [PATCH] done adding main logic in run_train_nerf to initialize processes with more robust handling of DDP --- mattport/nerf/trainer.py | 50 +------ scripts/run_train_nerf.py | 173 ++++++++++++++++++++++- tests/nerf/configs/test_basic_graph.yaml | 6 +- 3 files changed, 182 insertions(+), 47 deletions(-) diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 4bf367231..549934573 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -8,8 +8,6 @@ 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 @@ -54,52 +52,16 @@ class Trainer: """ raise NotImplementedError - 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): + def train(self) -> None: """_summary_""" raise NotImplementedError - def train_iteration(self, gpu: int): + def train_epoch(self): + """_summary_""" + raise NotImplementedError + + def train_iteration(self): """_summary_""" # TODO(): save checkpoints, and do logging only on rank 0 device only # if gpu == 0: 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 791020ac1..90a1ab246 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,21 +1,190 @@ """ 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.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(config) + @hydra.main(config_path="../configs", config_name="config.yaml") def main(config: DictConfig): """Main function.""" - trainer = Trainer(config) - trainer.run_multiprocess_train() + 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__": diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/configs/test_basic_graph.yaml index 77939a274..96d5c56ba 100644 --- a/tests/nerf/configs/test_basic_graph.yaml +++ b/tests/nerf/configs/test_basic_graph.yaml @@ -1,4 +1,8 @@ -world_size: 0 +machine_config: + num_gpus: 8 + num_machines: 1 + machine_rank: 0 + dist_url: auto network: mlp_0: class_name: MLP