diff --git a/docs/Trainer/Distributed training.md b/docs/Trainer/Distributed training.md index a5fdb24d..f40f5cbf 100644 --- a/docs/Trainer/Distributed training.md +++ b/docs/Trainer/Distributed training.md @@ -8,7 +8,16 @@ None of the flags below require changing anything about your lightningModel defi Lightning supports two backends. DataParallel and DistributedDataParallel. Both can be used for single-node multi-GPU training. For multi-node training you must use DistributedDataParallel. -**Warning: Your cluster must have NCCL installed and you must load it when submitting your SLURM script** +##### DataParallel (dp) +Splits a batch across multiple GPUs on the same node. Cannot be used for multi-node training. + +##### DistributedDataParallel (ddp) +Trains a copy of the model on each GPU and only syncs gradients. If used with DistributedSampler, each GPU trains +on a subset of the full dataset. + +##### DistributedDataParallel-2 (ddp2) +Works like DDP, except each node trains a single copy of the model using ALL GPUs on that node. +Very useful when dealing with negative samples, etc... You can toggle between each mode by setting this flag. ``` {.python} @@ -20,6 +29,9 @@ trainer = Trainer(distributed_backend='dp') # change to distributed data parallel (gpus > 1) trainer = Trainer(distributed_backend='ddp') + +# change to distributed data parallel (gpus > 1) +trainer = Trainer(distributed_backend='ddp2') ``` If you request multiple nodes, the back-end will auto-switch to ddp. @@ -144,11 +156,6 @@ script for the above trainer configuration. # activate conda env conda activate my_env -# REQUIRED: Load the latest NCCL version -# the nccl version must match the cuda used to build your PyTorch distribution -# (ie: which instructions did you follow when installing PyTorch) -# module load NCCL/2.4.7-1-cuda.10.0 - # ------------------------- # OPTIONAL # ------------------------- @@ -156,6 +163,10 @@ conda activate my_env # export NCCL_DEBUG=INFO # export PYTHONFAULTHANDLER=1 +# PyTorch comes with prebuilt NCCL support... but if you have issues with it +# you might need to load the latest version from your modules +# module load NCCL/2.4.7-1-cuda.10.0 + # on your cluster you might need these: # set the network interface # export NCCL_SOCKET_IFNAME=^docker0,lo diff --git a/examples/new_project_templates/lightning_module_template.py b/examples/new_project_templates/lightning_module_template.py index 3be61cf3..bfa65e5c 100644 --- a/examples/new_project_templates/lightning_module_template.py +++ b/examples/new_project_templates/lightning_module_template.py @@ -160,7 +160,7 @@ class LightningTemplateModel(LightningModule): # reduce manually when using dp val_acc = output['val_acc'] - if self.trainer.use_dp: + if self.trainer.use_dp or self.trainer.use_ddp2: val_acc = torch.mean(val_acc) val_acc_mean += val_acc diff --git a/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py b/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py index 7f201403..9367e9d0 100644 --- a/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py +++ b/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py @@ -107,7 +107,7 @@ def optimize_on_cluster(hyperparams): cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus cluster.per_experiment_nb_nodes = hyperparams.nb_gpu_nodes cluster.job_time = '2:00:00' - cluster.gpu_type = 'volta' + cluster.gpu_type = hyperparams.gpu_type cluster.memory_mb_per_node = 0 # any modules for code to run in env @@ -117,7 +117,7 @@ def optimize_on_cluster(hyperparams): cluster.add_command(f'export MASTER_PORT={PORT}') # OPTIONAL for debugging - # without these flags errors in your code will + # without these flags errors in your code will # appear to be nccl errors cluster.add_command('export NCCL_DEBUG=INFO') cluster.add_command('export PYTHONFAULTHANDLER=1') @@ -131,8 +131,6 @@ def optimize_on_cluster(hyperparams): # cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0']) # run only on 32GB voltas - cluster.add_slurm_cmd(cmd='constraint', value='volta32gb', - comment='use 32gb gpus') cluster.add_slurm_cmd(cmd='partition', value=hyperparams.gpu_partition, comment='your cluster might need this argument') @@ -181,6 +179,7 @@ if __name__ == '__main__': parent_parser.add_argument('--conda_env', type=str, default='base', help='email for jobs') parent_parser.add_argument('--gpu_partition', type=str, help='consult your cluster manual') + parent_parser.add_argument('--gpu_type', type=str, default='2080ti', help='consult your cluster manual') # allow model to overwrite or extend args parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir) @@ -191,4 +190,4 @@ if __name__ == '__main__': # --------------------- # run on HPC cluster print('RUNNING ON SLURM CLUSTER') - optimize_on_cluster(hyperparams) + optimize_on_cluster(hyperparams) \ No newline at end of file diff --git a/pytorch_lightning/pt_overrides/override_data_parallel.py b/pytorch_lightning/pt_overrides/override_data_parallel.py index b13b7559..dd3fe2d2 100644 --- a/pytorch_lightning/pt_overrides/override_data_parallel.py +++ b/pytorch_lightning/pt_overrides/override_data_parallel.py @@ -87,7 +87,6 @@ class LightningDistributedDataParallel(DistributedDataParallel): # -------------- # normal # output = self.module(*inputs[0], **kwargs[0]) - # lightning if self.module.training: output = self.module.training_step(*inputs[0], **kwargs[0]) @@ -99,6 +98,7 @@ class LightningDistributedDataParallel(DistributedDataParallel): outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs) output = self.gather(outputs, self.output_device) else: + # normal output = self.module(*inputs, **kwargs) if torch.is_grad_enabled(): @@ -171,6 +171,14 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): # pragma: n with lock: results[i] = e + # TODO: fix hack (maybe not a hack) + # make sure each module knows what training state it's in... + # fixes weird bug where copies are out of sync + root_m = modules[0] + for m in modules[1:]: + m.training = root_m.training + m.testing = root_m.testing + if len(modules) > 1: threads = [threading.Thread(target=_worker, args=(i, module, input, kwargs, device)) diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py index 3a636b04..73cea905 100644 --- a/pytorch_lightning/root_module/memory.py +++ b/pytorch_lightning/root_module/memory.py @@ -35,7 +35,7 @@ class ModelSummary(object): out_sizes = [] input_ = self.model.example_input_array - if self.model.on_gpu: + if self.model.use_ddp: input_ = input_.cuda(0) if self.model.trainer.use_amp: diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 6b8ed1ea..03fffe92 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -108,7 +108,7 @@ class Trainer(TrainerIO): :param val_check_interval: int. Check val this frequently within a train epoch :param log_save_interval: int. Writes logs to disk this often :param row_log_interval: int. How often to add logging rows - :param distributed_backend: str. dp, or ddp. + :param distributed_backend: str. Options: 'dp', 'ddp', 'ddp2'. :param use_amp: Bool. If true uses apex for 16bit precision :param print_nan_grads: Bool. Prints nan gradients :param print_weights_summary: Bool. Prints summary of weights @@ -172,8 +172,10 @@ class Trainer(TrainerIO): # distributed backend choice self.use_ddp = False + self.use_ddp2 = False self.use_dp = False self.single_gpu = False + self.distributed_backend = distributed_backend self.__set_distributed_mode(distributed_backend, nb_gpu_nodes) # init flags for SLURM+ddp to work @@ -285,6 +287,7 @@ class Trainer(TrainerIO): gpus = self.data_parallel_device_ids if gpus is None: return 0 + if type(gpus) is list: return len(gpus) if type(gpus) is int: @@ -307,6 +310,11 @@ class Trainer(TrainerIO): if distributed_backend is not None: self.use_dp = distributed_backend == 'dp' self.use_ddp = distributed_backend == 'ddp' + self.use_ddp2 = distributed_backend == 'ddp2' + + # disable single gpu when using ddp2 + if self.use_ddp2: + self.single_gpu = False # multiple GPU case elif self.num_gpus > 1: @@ -314,6 +322,7 @@ class Trainer(TrainerIO): # DP, DDP case self.use_dp = distributed_backend == 'dp' self.use_ddp = distributed_backend == 'ddp' + self.use_ddp2 = distributed_backend == 'ddp2' elif distributed_backend is None: m = 'When using multiple GPUs set ' \ @@ -483,7 +492,7 @@ class Trainer(TrainerIO): output = model(*args) return output - # CPU, single GPU + # single GPU if self.single_gpu: # for single GPU put inputs on gpu manually root_gpu = 0 @@ -492,6 +501,7 @@ class Trainer(TrainerIO): batch = self.transfer_batch_to_gpu(batch, root_gpu) args[0] = batch + # CPU if test: output = model.test_step(*args) else: @@ -547,10 +557,12 @@ class Trainer(TrainerIO): eval_results = {} - # give model a chance to do something with the outputs (and method defined) - model = self.__get_model() + # with a single dataloader don't pass an array if len(dataloaders) == 1: outputs = outputs[0] + + # give model a chance to do something with the outputs (and method defined) + model = self.__get_model() if test and self.__is_overriden('test_end'): eval_results = model.test_end(outputs) elif self.__is_overriden('validation_end'): @@ -656,7 +668,11 @@ class Trainer(TrainerIO): # when using multi-node or DDP within a node start each module in a separate process if self.use_ddp: - if self.is_slurm_managing_tasks: + if self.use_ddp2: + task = int(os.environ['SLURM_LOCALID']) + self.ddp_train(task, model) + + elif self.is_slurm_managing_tasks: task = int(os.environ['SLURM_LOCALID']) self.ddp_train(task, model) else: @@ -775,8 +791,13 @@ class Trainer(TrainerIO): self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0 # determine which process we are and world size - self.proc_rank = self.node_rank * self.num_gpus + gpu_nb - self.world_size = self.nb_gpu_nodes * self.num_gpus + if self.use_ddp: + self.proc_rank = self.node_rank * self.num_gpus + gpu_nb + self.world_size = self.nb_gpu_nodes * self.num_gpus + + elif self.use_ddp2: + self.proc_rank = self.node_rank + self.world_size = self.nb_gpu_nodes # let the exp know the rank to avoid overwriting logs if self.logger is not None: @@ -793,9 +814,19 @@ class Trainer(TrainerIO): # MODEL # copy model to each gpu - torch.cuda.set_device(gpu_nb) + if self.distributed_backend == 'ddp': + torch.cuda.set_device(gpu_nb) model.cuda(gpu_nb) + # set model properties before going into wrapper + model.trainer = self + model.on_gpu = self.on_gpu + model.use_dp = self.use_dp + model.use_ddp2 = self.use_ddp2 + model.use_ddp = self.use_ddp + model.use_amp = self.use_amp + model.testing = self.testing + # override root GPU self.root_gpu = gpu_nb @@ -808,8 +839,17 @@ class Trainer(TrainerIO): ) self.optimizers = optimizers - model = LightningDistributedDataParallel(model, device_ids=[gpu_nb], - find_unused_parameters=True) + # DDP2 uses all GPUs on the machine + if self.distributed_backend == 'ddp': + device_ids = [gpu_nb] + elif self.use_ddp2: + device_ids = None + + model = LightningDistributedDataParallel( + model, + device_ids=device_ids, + find_unused_parameters=True + ) # continue training routine self.__run_pretrain_routine(model) @@ -869,6 +909,7 @@ class Trainer(TrainerIO): ref_model.on_gpu = self.on_gpu ref_model.use_dp = self.use_dp ref_model.use_ddp = self.use_ddp + ref_model.use_ddp2 = self.use_ddp2 ref_model.use_amp = self.use_amp ref_model.testing = self.testing @@ -1138,7 +1179,7 @@ class Trainer(TrainerIO): progress_output = output['progress'] # reduce progress metrics for tqdm when using dp - if self.use_dp: + if self.use_dp or self.use_ddp2: nb_gpus = self.num_gpus progress_output = reduce_distributed_output(progress_output, nb_gpus) @@ -1162,7 +1203,7 @@ class Trainer(TrainerIO): ) # when using dp need to reduce the loss - if self.use_dp: + if self.use_dp or self.use_ddp2: loss = reduce_distributed_output(loss, self.num_gpus) return loss, model_specific_tqdm_metrics_dic @@ -1300,7 +1341,6 @@ class Trainer(TrainerIO): dataloaders, max_batches, test) - self.__add_tqdm_metrics(eval_out_metrics) # hook diff --git a/tests/debug.py b/tests/debug.py index db87f2ed..a80da5b6 100644 --- a/tests/debug.py +++ b/tests/debug.py @@ -59,18 +59,14 @@ class CoolModel(pl.LightningModule): return DataLoader(MNIST('path/to/save', train=False), batch_size=32) -def get_model(): +def get_model(use_test_model=False): # set up model with these hyperparams - root_dir = os.path.dirname(os.path.realpath(__file__)) - hparams = Namespace(**{'drop_prob': 0.2, - 'batch_size': 32, - 'in_features': 28 * 28, - 'learning_rate': 0.001 * 8, - 'optimizer_name': 'adam', - 'data_root': os.path.join(root_dir, 'mnist'), - 'out_features': 10, - 'hidden_dim': 1000}) - model = LightningTemplateModel(hparams) + hparams = get_hparams() + + if use_test_model: + model = LightningTestModel(hparams) + else: + model = LightningTemplateModel(hparams) return model, hparams @@ -114,7 +110,7 @@ def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningT trained_model = module_class.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=on_gpu, - map_location=map_location) + ) assert trained_model is not None, 'loading model failed' @@ -160,7 +156,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): result = trainer.fit(model) # correct result and ok accuracy - assert result == 1, 'amp + ddp model failed to complete' + assert result == 1, 'amp + ddp model failed sto complete' # test model loading pretrained_model = load_model(exp, save_dir, on_gpu) @@ -218,79 +214,23 @@ def main(): Make sure DDP + AMP continue training correctly :return: """ - hparams = get_hparams() - model = LightningTestModel(hparams) - + """ + Make sure DDP2 works + :return: + """ + os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) + model, hparams = get_model() trainer_options = dict( show_progress_bar=True, - max_nb_epochs=4, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, gpus=2, - distributed_backend='dp', + print_weights_summary=True, + distributed_backend='ddp2' ) - save_dir = init_save_dir() - - # exp file to get meta - exp = get_exp(False) - exp.argparse(hparams) - exp.save() - - # exp file to get weights - checkpoint = ModelCheckpoint(save_dir) - - # add these to the trainer options - trainer_options['experiment'] = exp - trainer_options['checkpoint_callback'] = checkpoint - - # fit model - trainer = Trainer(**trainer_options) - trainer.is_slurm_managing_tasks = True - result = trainer.fit(model) - - # track epoch before saving - real_global_epoch = trainer.current_epoch - - # correct result and ok accuracy - assert result == 1, 'amp + dp model failed to complete' - - # --------------------------- - # HPC LOAD/SAVE - # --------------------------- - # save - trainer.hpc_save(save_dir, exp) - - # init new trainer - new_exp = get_exp(False, version=exp.version) - trainer_options['experiment'] = new_exp - trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir) - trainer_options['train_percent_check'] = 0.2 - trainer_options['val_percent_check'] = 0.2 - trainer_options['max_nb_epochs'] = 1 - new_trainer = Trainer(**trainer_options) - - # set the epoch start hook so we can predict before the model does the full training - def assert_good_acc(): - assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0 - - # if model and state loaded correctly, predictions will be good even though we - # haven't trained with the new loaded model - dp_model = new_trainer.model - dp_model.eval() - - _ = [run_prediction(dataloader, dp_model, dp=True) for dataloader in trainer.val_dataloader] - - # new model - model = LightningTestModel(hparams) - model.on_sanity_check_start = assert_good_acc - - # fit new model which should load hpc weights - new_trainer.fit(model) - - # test freeze on gpu - model.freeze() - model.unfreeze() - - clear_save_dir() + run_gpu_model_test(trainer_options, model, hparams) if __name__ == '__main__': diff --git a/tests/test_models.py b/tests/test_models.py index 33abe410..62a7fef3 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -39,6 +39,29 @@ np.random.seed(SEED) # ------------------------------------------------------------------------ # TESTS # ------------------------------------------------------------------------ +def test_multi_gpu_model_ddp2(): + """ + Make sure DDP2 works + :return: + """ + if not can_run_gpu_test(): + return + + os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) + model, hparams = get_model() + trainer_options = dict( + show_progress_bar=True, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + gpus=2, + print_weights_summary=False, + distributed_backend='ddp2' + ) + + run_gpu_model_test(trainer_options, model, hparams) + + def test_dp_resume(): """ Make sure DP continues training correctly