release v0.3.4

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

pytorch_lightning/models/trainer.py

@@ -1,12 +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
@@ -146,17 +146,24 @@ 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"
@@ -170,6 +177,15 @@ class Trainer(TrainerIO):
             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 and self.use_dp:
+                self.use_ddp = True
+                self.use_dp = 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
 
@@ -372,15 +388,36 @@ class Trainer(TrainerIO):
 
         # when using multi-node or DDP within a node start each module in a separate process
         if self.use_ddp:
-            print('using 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.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(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:
-            print('using dp')
             self.dp_train(model)
 
         # ON CPU
@@ -416,7 +453,6 @@ class Trainer(TrainerIO):
             )
             self.optimizers = optimizers
 
-
         self.__run_pretrain_routine(model)
 
     def ddp_train(self, gpu_nb, model):
@@ -430,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
@@ -440,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
@@ -488,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):
         """
@@ -672,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):
@@ -724,6 +775,11 @@ class Trainer(TrainerIO):
         else:
             loss.backward()
 
+        # insert after step hook
+        if self.__is_function_implemented('on_after_backward'):
+            model_ref = self.__get_model()
+            response = model_ref.on_after_backward()
+
         if self.print_nan_grads:
             model = self.__get_model()
             for param in model.parameters():
@@ -744,6 +800,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/root_module/hooks.py

@@ -22,3 +22,24 @@ 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
+
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass
+

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.5',
+    version='0.3.4',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",