Ddp2 (#261)

* adds ddp2 option where on each node a single process uses all gpus * added ddp2 test * added ddp2 docs * Update Distributed training.md * delete ref to old update_training_log_metrics * delete ref to old update_training_log_metrics * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * banana pancakes * debug * debug * debug * debug * debug * debug * debug * debug * cheesecake
2026-09-12 12:40:20 +08:00 · 2019-10-04 15:07:54 -04:00
parent 2d335c664c
commit 32e74b8f36
8 changed files with 129 additions and 108 deletions
@@ -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
@@ -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
@@ -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)
@@ -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))
@@ -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:
@@ -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
@@ -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__':
@@ -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