release v0.3.3

docs/Trainer/Distributed training.md

@@ -3,6 +3,26 @@ Lightning makes multi-gpu training and 16 bit training trivial.
 *Note:*   
 None of the flags below require changing anything about your lightningModel definition. 
 
+---
+#### Choosing a backend  
+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.   
+
+You can toggle between each mode by setting this flag.
+``` {.python}
+# DEFAULT uses DataParallel
+trainer = Trainer(distributed_backend='dp')
+
+# change to distributed data parallel
+trainer = Trainer(distributed_backend='ddp')
+```
+
+If you request multiple nodes, the back-end will auto-switch to ddp.
+We recommend you use DistributedDataparallel even for single-node multi-GPU training. It is MUCH faster than DP but *may*
+have configuration issues depending on your cluster.
+
+For a deeper understanding of what lightning is doing, feel free to read [this guide](https://medium.com/@_willfalcon/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565).   
+
 ---
 #### 16-bit mixed precision
 16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.    
@@ -37,16 +57,69 @@ Make sure you're on a GPU machine. You can set as many GPUs as you want.
 In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
 ```python
 # set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+# lightning sets these flags for you automatically
+# no need to set yourself
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
 
-# DEFAULT
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+
+# to use DataParallel (default)
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+
+# RECOMMENDED use DistributedDataParallel
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
 ```
 
 ---
 #### Multi-node
-COMING SOON.
+Multi-node training is easily done by specifying these flags.
+```python
+# train on 12*8 GPUs
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
+```
+
+In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.
+
+```python
+cluster = SlurmCluster(
+    hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
+    log_path='/some/path/to/save',
+)
+
+# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
+# which interface your nodes use for communication
+cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+# see output of the NCCL connection process
+# NCCL is how the nodes talk to each other
+cluster.add_command('export NCCL_DEBUG=INFO')
+
+# setting a master port here is a good idea.
+cluster.add_command(f'export MASTER_PORT={PORT}')
+
+# good to load the latest NCCL version
+cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+
+# configure cluster
+cluster.per_experiment_nb_nodes = 12 
+cluster.per_experiment_nb_gpus = 8
+
+cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
+```
+
+Finally, make sure to add a distributed sampler to your dataset. The distributed sampler copies a 
+portion of your dataset onto each GPU. (World_size = gpus_per_node * nb_nodes).   
+
+```python
+# ie: this:
+dataset = myDataset()
+dataloader = Dataloader(dataset)
+
+# becomes:
+dataset = myDataset()
+dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+dataloader = Dataloader(dataset, sampler=dist_sampler)
+```
 
 ---
 #### Self-balancing architecture

docs/Trainer/Training Loop.md

@@ -19,7 +19,7 @@ Cut the learning rate by 10 at every epoch listed in this list.
 trainer = Trainer(lr_scheduler_milestones=None)
 
 # cut LR by 10 at 100, 200, and 300 epochs 
-trainer = Trainer(lr_scheduler_milestones=[100, 200, 300])
+trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
 ```
 
 ---

examples/new_project_templates/lightning_module_template.py

@@ -84,9 +84,10 @@ class LightningTemplateModel(LightningModule):
         loss_val = self.loss(y, y_hat)
 
         output = OrderedDict({
-            'loss': loss_val,
-            'tqdm_metrics': {}
+            'loss': loss_val
         })
+
+        # can also return just a scalar instead of a dict (return loss_val)
         return output
 
     def validation_step(self, data_batch, batch_i):
@@ -107,8 +108,10 @@ class LightningTemplateModel(LightningModule):
 
         output = OrderedDict({
             'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc),
+            'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index),
         })
+
+        # can also return just a scalar instead of a dict (return loss_val)
         return output
 
     def validation_end(self, outputs):
@@ -117,6 +120,10 @@ class LightningTemplateModel(LightningModule):
         :param outputs: list of individual outputs of each validation step
         :return:
         """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
         val_loss_mean = 0
         val_acc_mean = 0
         for output in outputs:

examples/new_project_templates/single_gpu_node_dp_template.py

@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

examples/new_project_templates/single_gpu_template.py

@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='0', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

pytorch_lightning/models/trainer.py

@@ -1,11 +1,12 @@
 """
 The trainer handles all the logic for running a val loop, training loop, distributing, etc...
 """
-from time import sleep
 import subprocess
 import traceback
 import warnings
 import os
+import pdb
+import re
 
 import torch
 from torch.utils.data.distributed import DistributedSampler
@@ -17,7 +18,7 @@ import tqdm
 
 from pytorch_lightning.root_module.memory import get_gpu_memory_map
 from pytorch_lightning.root_module.model_saving import TrainerIO
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
 
 
 try:
@@ -27,6 +28,27 @@ except ModuleNotFoundError:
     APEX_AVAILABLE = False
 
 
+def reduce_distributed_output(output, nb_gpus):
+    if nb_gpus <= 1:
+        return output
+
+    # when using DP, we get one output per gpu
+    # average outputs and return
+    if type(output) is torch.Tensor:
+        return output.mean()
+
+    for k, v in output.items():
+        # recurse on nested dics
+        if isinstance(output[k], dict):
+            output[k] = reduce_distributed_output(output[k], nb_gpus)
+
+        # reduce only metrics that have the same nb of gpus
+        elif output[k].size(0) == nb_gpus:
+            reduced = torch.mean(output[k])
+            output[k] = reduced
+    return output
+
+
 class Trainer(TrainerIO):
 
     def __init__(self,
@@ -50,12 +72,47 @@ class Trainer(TrainerIO):
                  val_check_interval=0.95,
                  log_save_interval=100, add_log_row_interval=10,
                  lr_scheduler_milestones=None,
+                 distributed_backend='dp',
                  use_amp=False,
                  print_nan_grads=False,
                  print_weights_summary=True,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
 
+        """
+
+        :param experiment: Test-tube experiment
+        :param early_stop_callback: from pytorch_lightning import EarlyStopping
+        :param checkpoint_callback: from pytorch_lightning import Checkpoint
+        :param gradient_clip:
+        :param cluster:
+        :param process_position:
+        :param current_gpu_name:
+        :param nb_gpu_nodes:
+        :param gpus:
+        :param progress_bar:
+        :param overfit_pct:
+        :param track_grad_norm:
+        :param check_val_every_n_epoch:
+        :param fast_dev_run:
+        :param accumulate_grad_batches:
+        :param max_nb_epochs:
+        :param min_nb_epochs:
+        :param train_percent_check:
+        :param val_percent_check:
+        :param test_percent_check:
+        :param val_check_interval:
+        :param log_save_interval:
+        :param add_log_row_interval:
+        :param lr_scheduler_milestones:
+        :param distributed_backend: 'np' to use DistributedParallel, 'ddp' to use DistributedDataParallel
+        :param use_amp:
+        :param print_nan_grads:
+        :param print_weights_summary:
+        :param amp_level:
+        :param nb_sanity_val_steps:
+        """
+
         # Transfer params
         self.nb_gpu_nodes = nb_gpu_nodes
         self.gradient_clip = gradient_clip
@@ -89,21 +146,45 @@ class Trainer(TrainerIO):
         self.print_nan_grads = print_nan_grads
         self.data_parallel_device_ids = None
         self.world_size = 1
+        self.node_rank = 0
+        self.use_ddp = False
+        self.use_dp = False
+
 
         # gpus come in as a string.
         # if gpus = -1 then use all available devices
         # otherwise, split the string using commas
         if gpus is not None:
-            if gpus == '-1':
-                self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+            if type(gpus) is list:
+                self.data_parallel_device_ids = gpus
+            elif type(gpus) is str:
+                if gpus == '-1':
+                    self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+                else:
+                    self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
             else:
-                self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+                raise Exception('gpus has to be a string or list of ids')
 
             # set the correct cuda visible devices (using pci order)
             os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
             os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+            print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
 
-        self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 1
+        # make DP and DDP mutually exclusive
+        # single GPU will also use DP with devices=[0]
+        have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+        if have_gpus:
+            self.use_dp = distributed_backend == 'dp'
+            self.use_ddp = distributed_backend == 'ddp'
+
+            # use ddp automatically if nb_gpu_nodes > 1
+            if nb_gpu_nodes > 1:
+                self.use_ddp = True
+                self.use_ddp = False
+                w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+                    'Switching to DistributedDataParallel for you. ' \
+                    'To silence this warning set distributed_backend=ddp'
+                warnings.warn(w)
 
         # process info
         self.proc_rank = 0
@@ -144,6 +225,10 @@ class Trainer(TrainerIO):
             '''
             warnings.warn(msg)
 
+    @property
+    def data_parallel(self):
+        return self.use_dp or self.use_ddp
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -242,8 +327,12 @@ class Trainer(TrainerIO):
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            if self.data_parallel:
+            if self.use_ddp:
                 output = model(data_batch, batch_i)
+            elif self.use_dp:
+                output = model(data_batch, batch_i)
+                output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+
             else:
                 output = model.validation_step(data_batch, batch_i)
 
@@ -277,7 +366,7 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
-        if self.data_parallel and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
+        if self.use_ddp and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
             msg = '''
             when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
             
@@ -297,16 +386,41 @@ class Trainer(TrainerIO):
     # -----------------------------
     def fit(self, model):
 
-        # when using gpus, first thing we do is spawn a new process between each worker
-        #  multi-gpu and multi-nodes
-        if self.data_parallel:
+        # when using multi-node or DDP within a node start each module in a separate process
+        if self.use_ddp:
+            # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
             self.experiment = self.experiment.get_meta_copy()
-            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
 
-        # treat 1 gpu as a different case to avoid nccl bugs
-        elif self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) == 1:
-            self.single_gpu_train(model)
+            # whenever we have the correct number of tasks, we let slurm manage processes
+            # otherwise we launch the required number of processes
+            nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+            nb_slurm_tasks = 0
+            try:
+                nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
+            except Exception as e:
+                # likely not on slurm, so set the slurm managed flag to false
+                is_slurm_managing_tasks = False
 
+            if is_slurm_managing_tasks:
+                task = int(os.environ['SLURM_LOCALID'])
+                self.ddp_train(task, model)
+            else:
+                msg = f"""
+                You requested {nb_requested_gpus} GPUs but launched {nb_slurm_tasks} slurm tasks. 
+                We will launch {nb_requested_gpus} processes for you. 
+                We recommend you let slurm manage the processes by setting: --ntasks-per-node={nb_requested_gpus}
+                If you're not using SLURM, ignore this message!
+                """
+                warnings.warn(msg)
+                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+
+        # 1 gpu or dp option triggers training using DP module
+        # easier to avoid NCCL issues
+        elif self.use_dp:
+            self.dp_train(model)
+
+        # ON CPU
         else:
             # CHOOSE OPTIMIZER
             # filter out the weights that were done on gpu so we can load on good old cpus
@@ -322,14 +436,15 @@ class Trainer(TrainerIO):
 
             self.__run_pretrain_routine(model)
 
-    def single_gpu_train(self, model):
-        # torch.cuda.set_device(0)
-        model.cuda(0)
+    def dp_train(self, model):
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        model.cuda(self.data_parallel_device_ids[0])
+        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
         # run through amp wrapper
         if self.use_amp:
             # An example
@@ -340,7 +455,7 @@ class Trainer(TrainerIO):
 
         self.__run_pretrain_routine(model)
 
-    def dp_train(self, gpu_nb, model):
+    def ddp_train(self, gpu_nb, model):
         """
         Entry point into a DP thread
         :param gpu_nb:
@@ -351,9 +466,10 @@ class Trainer(TrainerIO):
         # node rank using relative slurm id
         # otherwise default to node rank 0
         try:
-            node_rank = int(os.environ['SLURM_NODEID'])
-        except KeyError as e:
-            node_rank = 0
+            node_id = os.environ['SLURM_NODEID']
+            self.node_rank = int(node_id)
+        except Exception as e:
+            self.node_rank = 0
 
         # recover original exp before went into process
         # init in write mode only on proc 0
@@ -361,10 +477,10 @@ class Trainer(TrainerIO):
         self.experiment = self.experiment.get_non_ddp_exp()
 
         # show progbar only on prog_rank 0
-        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
+        self.prog_bar = self.prog_bar and self.node_rank == 0 and gpu_nb == 0
 
         # determine which process we are and world size
-        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
         self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
         # set up server using proc 0's ip address
@@ -409,15 +525,29 @@ class Trainer(TrainerIO):
             port = 12910
             os.environ['MASTER_PORT'] = f'{port}'
 
+        root_node = self.__resolve_root_node_address()
+        os.environ['MASTER_ADDR'] = root_node
+
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+
+    def __resolve_root_node_address(self):
         try:
             root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+
+            if '[' in root_node:
+                name = root_node.split('[')[0]
+                number = root_node.split(',')[0]
+                if '-' in number:
+                    number = number.split('-')[0]
+
+                number = re.sub('[^0-9]', '', number)
+                root_node = name + number
+
         except Exception as e:
             root_node = '127.0.0.2'
 
-        os.environ['MASTER_ADDR'] = root_node
-
-        sleep(self.proc_rank*0.5)
-        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+        return root_node
 
     def __run_pretrain_routine(self, model):
         """
@@ -593,7 +723,7 @@ class Trainer(TrainerIO):
 
     def __log_vals_blacklist(self):
         """avoid logging some vals lightning uses to maintain state"""
-        blacklist = {'batch_nb', 'v_nb', 'epoch', 'gpu'}
+        blacklist = {'batch_nb', 'v_nb', 'gpu'}
         return blacklist
 
     def __run_tng_batch(self, data_batch, batch_nb):
@@ -602,8 +732,8 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            model = self.__get_model()
-            response = model.on_batch_start(data_batch)
+            model_ref = self.__get_model()
+            response = model_ref.on_batch_start(data_batch)
 
             if response == -1:
                 return -1
@@ -613,8 +743,11 @@ class Trainer(TrainerIO):
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        if self.data_parallel:
+        if self.use_ddp:
             output = self.model(data_batch, batch_nb)
+        elif self.use_dp:
+            output = self.model(data_batch, batch_nb)
+            output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
         else:
             output = self.model.training_step(data_batch, batch_nb)
 
@@ -662,6 +795,11 @@ class Trainer(TrainerIO):
             for optimizer in self.optimizers:
                 optimizer.step()
 
+                # insert after step hook
+                if self.__is_function_implemented('on_before_zero_grad'):
+                    model_ref = self.__get_model()
+                    response = model_ref.on_before_zero_grad(optimizer)
+
                 # clear gradients
                 optimizer.zero_grad()
 

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -1,6 +1,7 @@
 from torch.nn import DataParallel
 from torch.nn.parallel import DistributedDataParallel
 import itertools
+from itertools import chain
 
 import threading
 import torch
@@ -42,6 +43,29 @@ class LightningDataParallel(DataParallel):
     Override the forward call in lightning so it goes to training and validation step respectively
     """
 
+    def forward(self, *inputs, **kwargs):
+        if not self.device_ids:
+            return self.module(*inputs, **kwargs)
+
+        for t in chain(self.module.parameters(), self.module.buffers()):
+            if t.device != self.src_device_obj:
+                raise RuntimeError("module must have its parameters and buffers "
+                                   "on device {} (device_ids[0]) but found one of "
+                                   "them on device: {}".format(self.src_device_obj, t.device))
+
+        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+        if len(self.device_ids) == 1:
+            # lightning
+            if self.module.training:
+                return self.module.training_step(*inputs[0], **kwargs[0])
+            else:
+                return self.module.validation_step(*inputs[0], **kwargs[0])
+
+        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
+        outputs = self.parallel_apply(replicas, inputs, kwargs)
+        return self.gather(outputs, self.output_device)
+
+
     def parallel_apply(self, replicas, inputs, kwargs):
         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
 

pytorch_lightning/root_module/hooks.py

@@ -22,3 +22,17 @@ class ModelHooks(torch.nn.Module):
     def on_tng_metrics(self, metrics):
         pass
 
+    def on_before_zero_grad(self, optimizer):
+        """
+        Called after optimizer.step() and before optimizer.zero_grad()
+
+        for optimizer in optimizers:
+            optimizer.step()
+            model.on_before_zero_grad(optimizer) # < ---- called here
+            optimizer.zero_grad
+
+        :param optimizer:
+        :return:
+        """
+        pass
+

pytorch_lightning/root_module/model_saving.py

@@ -2,7 +2,7 @@ import torch
 import os
 import re
 import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
 
 class ModelIO(object):
 
@@ -66,7 +66,8 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
+        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+        model = self.model.module if is_dp_module else self.model
         checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.2.4.1',
+    version='0.3.3',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",