Repository: Lightning-AI/pytorch-lightning

Issue:
Trainer DDP invoking load_spawn_weights() on each node
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
On a SLURM cluster, I am seeing the same problem as issue #1335 , despite that issue's fix being applied.  

### To Reproduce

Steps to reproduce the behavior:
1. Allocate 4 nodes on a SLURM-managed cluster
1. srun the script pl.py on each allocated node
2. See the errors on 3 of 4 nodes:
```
INFO:lightning:GPU available: True, used: True
INFO:lightning:VISIBLE GPUS: 0,1
INFO:lightning:GPU available: True, used: True
INFO:lightning:VISIBLE GPUS: 0,1
INFO:lightning:GPU available: True, used: True
INFO:lightning:VISIBLE GPUS: 0,1
INFO:lightning:GPU available: True, used: True
INFO:lightning:VISIBLE GPUS: 0,1
/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: You have defined a `test_dataloader()` and have defined a `test_step()`, you may also want to define `test_epoch_end()` for accumulating stats.
  warnings.warn(*args, **kwargs)
/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: You have defined a `test_dataloader()` and have defined a `test_step()`, you may also want to define `test_epoch_end()` for accumulating stats.
  warnings.warn(*args, **kwargs)
/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: You have defined a `test_dataloader()` and have defined a `test_step()`, you may also want to define `test_epoch_end()` for accumulating stats.
  warnings.warn(*args, **kwargs)
/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: You have defined a `test_dataloader()` and have defined a `test_step()`, you may also want to define `test_epoch_end()` for accumulating stats.
  warnings.warn(*args, **kwargs)
d-12-3-1:47016:47016 [0] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.132<0>
d-12-3-1:47016:47016 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-3-1:47016:47016 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.132<0>
NCCL version 2.4.8+cuda10.1
d-12-3-1:47022:47022 [1] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.132<0>
d-12-3-1:47022:47022 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-4-1:51815:51815 [0] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.134<0>
d-12-4-1:51820:51820 [1] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.134<0>
d-12-4-1:51820:51820 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-4-1:51815:51815 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-4-2:51993:51993 [0] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.135<0>
d-12-4-2:51997:51997 [1] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.135<0>
d-12-3-2:43991:43991 [1] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.133<0>
d-12-3-2:43985:43985 [0] NCCL INFO Bootstrap : Using [0]ens2f0:172.31.130.133<0>
d-12-4-2:51993:51993 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-4-2:51997:51997 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-3-2:43991:43991 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-3-2:43985:43985 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so).
d-12-3-1:47022:47022 [1] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.132<0>
d-12-3-1:47022:47036 [1] NCCL INFO Setting affinity for GPU 1 to ffff,f00000ff,fff00000
d-12-4-2:51993:51993 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.135<0>
d-12-4-2:51997:51997 [1] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.135<0>
d-12-4-1:51820:51820 [1] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.134<0>
d-12-3-2:43991:43991 [1] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.133<0>
d-12-3-2:43985:43985 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.133<0>
d-12-4-1:51815:51815 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/RoCE ; OOB ens2f0:172.31.130.134<0>
d-12-4-2:51993:52008 [0] NCCL INFO Setting affinity for GPU 0 to ffff,f00000ff,fff00000
d-12-4-2:51997:52010 [1] NCCL INFO Setting affinity for GPU 1 to ffff,f00000ff,fff00000
d-12-4-1:51820:51830 [1] NCCL INFO Setting affinity for GPU 1 to ffff,f00000ff,fff00000
d-12-3-2:43991:44003 [1] NCCL INFO Setting affinity for GPU 1 to ffff,f00000ff,fff00000
d-12-3-2:43985:44002 [0] NCCL INFO Setting affinity for GPU 0 to ffff,f00000ff,fff00000
d-12-4-1:51815:51832 [0] NCCL INFO Setting affinity for GPU 0 to ffff,f00000ff,fff00000
d-12-3-1:47016:47034 [0] NCCL INFO Setting affinity for GPU 0 to ffff,f00000ff,fff00000
d-12-4-2:51997:52010 [1] NCCL INFO CUDA Dev 1[1], IB NIC distance :  NODE
d-12-4-2:51993:52008 [0] NCCL INFO CUDA Dev 0[0], IB NIC distance :  NODE
d-12-4-1:51815:51832 [0] NCCL INFO CUDA Dev 0[0], IB NIC distance :  NODE
d-12-4-1:51820:51830 [1] NCCL INFO CUDA Dev 1[1], IB NIC distance :  NODE
d-12-3-2:43991:44003 [1] NCCL INFO CUDA Dev 1[1], IB NIC distance :  NODE
d-12-3-2:43985:44002 [0] NCCL INFO CUDA Dev 0[0], IB NIC distance :  NODE
d-12-3-1:47022:47036 [1] NCCL INFO CUDA Dev 1[1], IB NIC distance :  NODE
d-12-3-1:47016:47034 [0] NCCL INFO CUDA Dev 0[0], IB NIC distance :  NODE
d-12-3-1:47016:47034 [0] NCCL INFO Channel 00 :    0   1   2   3   4   5   6   7
d-12-3-1:47016:47034 [0] NCCL INFO Ring 00 : 7 -> 0 [receive] via NET/IB/0
d-12-4-2:51993:52008 [0] NCCL INFO Ring 00 : 5 -> 6 [receive] via NET/IB/0
d-12-3-2:43985:44002 [0] NCCL INFO Ring 00 : 1 -> 2 [receive] via NET/IB/0
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 3 -> 4 [receive] via NET/IB/0
d-12-3-2:43985:44002 [0] NCCL INFO Ring 00 : 2[0] -> 3[1] via direct shared memory
d-12-4-2:51993:52008 [0] NCCL INFO Ring 00 : 6[0] -> 7[1] via direct shared memory
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 4[0] -> 5[1] via direct shared memory
d-12-3-1:47016:47034 [0] NCCL INFO Ring 00 : 0[0] -> 1[1] via direct shared memory
d-12-4-2:51997:52010 [1] NCCL INFO Ring 00 : 7 -> 0 [send] via NET/IB/0
d-12-3-2:43991:44003 [1] NCCL INFO Ring 00 : 3 -> 4 [send] via NET/IB/0
d-12-4-1:51820:51830 [1] NCCL INFO Ring 00 : 5 -> 6 [send] via NET/IB/0
d-12-3-1:47022:47036 [1] NCCL INFO Ring 00 : 1 -> 2 [send] via NET/IB/0
d-12-3-2:43991:44003 [1] NCCL INFO Ring 00 : 3[1] -> 2[0] via direct shared memory
d-12-4-1:51820:51830 [1] NCCL INFO Ring 00 : 5[1] -> 4[0] via direct shared memory
d-12-4-2:51997:52010 [1] NCCL INFO Ring 00 : 7[1] -> 6[0] via direct shared memory
d-12-3-1:47022:47036 [1] NCCL INFO Ring 00 : 1[1] -> 0[0] via direct shared memory
d-12-4-1:51820:51830 [1] NCCL INFO Trees [0] 4->5->-1/-1/-1
d-12-4-2:51997:52010 [1] NCCL INFO Trees [0] 6->7->-1/-1/-1
d-12-3-2:43991:44003 [1] NCCL INFO Trees [0] 2->3->-1/-1/-1
d-12-3-1:47022:47036 [1] NCCL INFO Trees [0] 0->1->-1/-1/-1
d-12-4-1:51820:51830 [1] NCCL INFO comm 0x7fa5380022f0 rank 5 nranks 8 cudaDev 1 nvmlDev 1 - Init COMPLETE
d-12-3-2:43991:44003 [1] NCCL INFO comm 0x7ff3f40022f0 rank 3 nranks 8 cudaDev 1 nvmlDev 1 - Init COMPLETE
d-12-4-2:51997:52010 [1] NCCL INFO comm 0x7f84600022f0 rank 7 nranks 8 cudaDev 1 nvmlDev 1 - Init COMPLETE
d-12-3-1:47022:47036 [1] NCCL INFO comm 0x7f1f780022f0 rank 1 nranks 8 cudaDev 1 nvmlDev 1 - Init COMPLETE
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 2 -> 4 [receive] via NET/IB/0
d-12-3-2:43985:44002 [0] NCCL INFO Ring 00 : 2 -> 4 [send] via NET/IB/0
d-12-4-2:51993:52008 [0] NCCL INFO Ring 00 : 6 -> 4 [send] via NET/IB/0
d-12-3-1:47016:47034 [0] NCCL INFO Ring 00 : 4 -> 0 [receive] via NET/IB/0
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 6 -> 4 [receive] via NET/IB/0
d-12-3-2:43985:44002 [0] NCCL INFO Ring 00 : 4 -> 2 [receive] via NET/IB/0
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 4 -> 0 [send] via NET/IB/0
d-12-4-2:51993:52008 [0] NCCL INFO Ring 00 : 4 -> 6 [receive] via NET/IB/0
d-12-3-1:47016:47034 [0] NCCL INFO Ring 00 : 0 -> 4 [send] via NET/IB/0
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 0 -> 4 [receive] via NET/IB/0
d-12-3-1:47016:47034 [0] NCCL INFO Trees [0] -1->0->1/4/-1
d-12-3-1:47016:47034 [0] NCCL INFO Using 256 threads, Min Comp Cap 7, Trees enabled up to size 79999
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 4 -> 2 [send] via NET/IB/0
d-12-3-1:47016:47034 [0] NCCL INFO comm 0x7fa36c0022f0 rank 0 nranks 8 cudaDev 0 nvmlDev 0 - Init COMPLETE
d-12-3-1:47016:47016 [0] NCCL INFO Launch mode Parallel
d-12-4-1:51815:51832 [0] NCCL INFO Ring 00 : 4 -> 6 [send] via NET/IB/0
d-12-3-2:43985:44002 [0] NCCL INFO Trees [0] 4->2->3/-1/-1
d-12-3-2:43985:44002 [0] NCCL INFO comm 0x7febb00022f0 rank 2 nranks 8 cudaDev 0 nvmlDev 0 - Init COMPLETE
d-12-4-1:51815:51832 [0] NCCL INFO Trees [0] 0->4->5/2/6
d-12-4-2:51993:52008 [0] NCCL INFO Trees [0] 4->6->7/-1/-1
d-12-4-1:51815:51832 [0] NCCL INFO comm 0x7fab7c0022f0 rank 4 nranks 8 cudaDev 0 nvmlDev 0 - Init COMPLETE
d-12-4-2:51993:52008 [0] NCCL INFO comm 0x7f2bec0022f0 rank 6 nranks 8 cudaDev 0 nvmlDev 0 - Init COMPLETE
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
INFO:lightning:
  | Name    | Type   | Params
-------------------------------
0 | layer_1 | Linear | 100 K 
1 | layer_2 | Linear | 33 K  
2 | layer_3 | Linear | 2 K   
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
INFO:lightning:Set SLURM handle signals.
spr= 5 snr= 2 sng= 2 gpu_idx= 1
rank= 5  world_size= 8
Root node= d-12-3-1
spr= 6 snr= 3 sng= 2 gpu_idx= 0
rank= 6  world_size= 8
Root node= d-12-3-1
spr= 0 snr= 0 sng= 2 gpu_idx= 0
rank= 0  world_size= 8
Root node= d-12-3-1
spr= 3 snr= 1 sng= 2 gpu_idx= 1
rank= 3  world_size= 8
Root node= d-12-3-1
spr= 4 snr= 2 sng= 2 gpu_idx= 0
rank= 4  world_size= 8
Root node= d-12-3-1
spr= 1 snr= 0 sng= 2 gpu_idx= 1
rank= 1  world_size= 8
Root node= d-12-3-1
spr= 2 snr= 1 sng= 2 gpu_idx= 0
rank= 2  world_size= 8
Root node= d-12-3-1
spr= 7 snr= 3 sng= 2 gpu_idx= 1
rank= 7  world_size= 8
Root node= d-12-3-1
/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: Displayed epoch numbers in the progress bar start from "1" until v0.6.x, but will start from "0" in v0.8.0.
  warnings.warn(*args, **kwargs)
Epoch 2: 100%|##########| 118/118 [00:19<00:00,  5.98it/s, loss=0.197, v_num=436408]
before lsw: self.proc_rank= 0                             
load_spawn_weights called for self.proc_rank= 0
before lsw: self.proc_rank= 0
load_spawn_weights called for self.proc_rank= 0
Traceback (most recent call last):
  File "pl.py", line 119, in <module>
    main()
  File "pl.py", line 111, in main
    trainer.fit(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 697, in fit
    self.load_spawn_weights(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 376, in load_spawn_weights
    loaded_model = original_model.__class__.load_from_checkpoint(path)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/core/lightning.py", line 1504, in load_from_checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 525, in load
    with _open_file_like(f, 'rb') as opened_file:
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 212, in _open_file_like
    return _open_file(name_or_buffer, mode)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 193, in __init__
    super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/home/gridsan/groups/anaconda/dpc/lightning/__temp_weight_ddp_end.ckpt'
before lsw: self.proc_rank= 0
load_spawn_weights called for self.proc_rank= 0
Traceback (most recent call last):
  File "pl.py", line 119, in <module>
    main()
  File "pl.py", line 111, in main
    trainer.fit(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 697, in fit
    self.load_spawn_weights(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 376, in load_spawn_weights
    loaded_model = original_model.__class__.load_from_checkpoint(path)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/core/lightning.py", line 1504, in load_from_checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 525, in load
    with _open_file_like(f, 'rb') as opened_file:
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 212, in _open_file_like
    return _open_file(name_or_buffer, mode)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 193, in __init__
    super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/home/gridsan/groups/anaconda/dpc/lightning/__temp_weight_ddp_end.ckpt'
before lsw: self.proc_rank= 0
load_spawn_weights called for self.proc_rank= 0
Traceback (most recent call last):
  File "pl.py", line 119, in <module>
    main()
  File "pl.py", line 111, in main
    trainer.fit(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 697, in fit
    self.load_spawn_weights(model)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 376, in load_spawn_weights
    loaded_model = original_model.__class__.load_from_checkpoint(path)
  File "/home/gridsan/dcampbell/.local/lib/python3.6/site-packages/pytorch_lightning/core/lightning.py", line 1504, in load_from_checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 525, in load
    with _open_file_like(f, 'rb') as opened_file:
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 212, in _open_file_like
    return _open_file(name_or_buffer, mode)
  File "/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.6/site-packages/torch/serialization.py", line 193, in __init__
    super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/home/gridsan/groups/anaconda/dpc/lightning/__temp_weight_ddp_end.ckpt'
srun: error: d-12-4-1: task 2: Exited with exit code 1
srun: error: d-12-4-2: task 3: Exited with exit code 1
srun: error: d-12-3-1: task 0: Exited with exit code 1
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample

sbatch script:
```
#!/bin/bash -l

# SLURM SUBMIT SCRIPT
#SBATCH --nodes=4
#SBATCH --gres=gpu:volta:2
#SBATCH --mem=0
#SBATCH --time=0-02:00:00
#SBATCH --partition=gaia

# activate conda env
#source activate $1

# -------------------------
# debugging flags (optional)
export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1

# on your cluster you might need these:
# set the network interface
# export NCCL_SOCKET_IFNAME=^docker0,lo

# might need the latest cuda
# module load NCCL/2.4.7-1-cuda.10.0
# -------------------------

module load mpi/openmpi-4.0
module load anaconda/2020b

export MASTER_PORT=`comm -23 <(seq 12000 18000 | sort) <(ss -Htan | awk '{print $4}' | cut -d':' -f2 | sort -u) | shuf | head -n 1`

# run script from above
srun python pl.py
```

pl.py (dervied from lightning tutorial's MNIST example)
```
import torch
from torch import nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, random_split
from torch.nn import functional as F
from torchvision.datasets import MNIST
from torchvision import datasets, transforms
import os

import torch.multiprocessing as mp

from localtools import slurm_torch
from hostlist import expand_hostlist

class LightningMNISTClassifier(pl.LightningModule):

  def __init__(self):
    super(LightningMNISTClassifier, self).__init__()

    # mnist images are (1, 28, 28) (channels, width, height) 
    self.layer_1 = torch.nn.Linear(28 * 28, 128)
    self.layer_2 = torch.nn.Linear(128, 256)
    self.layer_3 = torch.nn.Linear(256, 10)

  def forward(self, x):
      batch_size, channels, width, height = x.size()

      # (b, 1, 28, 28) -> (b, 1*28*28)
      x = x.view(batch_size, -1)

      # layer 1 (b, 1*28*28) -> (b, 128)
      x = self.layer_1(x)
      x = torch.relu(x)

      # layer 2 (b, 128) -> (b, 256)
      x = self.layer_2(x)
      x = torch.relu(x)

      # layer 3 (b, 256) -> (b, 10)
      x = self.layer_3(x)

      # probability distribution over labels
      x = torch.log_softmax(x, dim=1)

      return x

  def cross_entropy_loss(self, logits, labels):
    return F.nll_loss(logits, labels)

  def training_step(self, train_batch, batch_idx):
      x, y = train_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)

      logs = {'train_loss': loss}
      return {'loss': loss, 'log': logs}

  def test_step(self, test_batch, batch_idx):
      x, y = test_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      return {'test_loss': loss}

  def validation_step(self, val_batch, batch_idx):
      x, y = val_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      return {'val_loss': loss}

  def validation_epoch_end(self, outputs):
      # called at the end of the validation epoch
      # outputs is an array with what you returned in validation_step for each batch
      # outputs = [{'loss': batch_0_loss}, {'loss': batch_1_loss}, ..., {'loss': batch_n_loss}] 
      avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
      tensorboard_logs = {'val_loss': avg_loss}
      return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

  def prepare_data(self):
    # transforms for images
    transform=transforms.Compose([transforms.ToTensor(), 
                                  transforms.Normalize((0.1307,), (0.3081,))])
      
    # prepare transforms standard to MNIST
    mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
    mnist_test = MNIST(os.getcwd(), train=False, download=True, transform=transform)
    
    self.mnist_train, self.mnist_val = random_split(mnist_train, [55000, 5000])

  def train_dataloader(self):
    return DataLoader(self.mnist_train, num_workers=16, batch_size=64)

  def val_dataloader(self):
    return DataLoader(self.mnist_val, num_workers=16, batch_size=64)

  def test_dataloader(self):
    return DataLoader(self,mnist_test, num_workers=16, batch_size=64)

  def configure_optimizers(self):
    optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
    return optimizer

# train
def main ():
    rank, size = slurm_torch.torch_env()
    nn = ' '.join(expand_hostlist( os.environ['SLURM_NODELIST']))
    os.environ['SLURM_NODELIST']=nn
    model = LightningMNISTClassifier()
    trainer = pl.Trainer(gpus=2, num_nodes=size, distributed_backend='ddp', max_epochs=2)
    trainer.fit(model)

if __name__ == '__main__':
    root_dir = os.path.dirname(os.path.realpath(__file__))

    # TRAIN
    main()
```
### Expected behavior

Workers on all nodes run to completion without errors

### Environment
```
* CUDA:
        - GPU:
                - Tesla V100-PCIE-32GB
                - Tesla V100-PCIE-32GB
        - available:         True
        - version:           10.1
* Packages:
        - numpy:             1.18.1
        - pyTorch_debug:     False
        - pyTorch_version:   1.4.0
        - pytorch-lightning: 0.7.3
        - tensorboard:       2.1.0
        - tqdm:              4.45.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - 
        - processor:         x86_64
        - python:            3.6.10
        - version:           #1 SMP Fri Apr 3 11:13:11 EDT 2020
```

 - How you installed PyTorch: pip


### Additional context

I'm not positive that this isn't a user-introduced problem as I had to do a little bit of tweaking to the supplied example application in order to run in my environment.  

I have added some outputs to attempt to determine what's going on, and it seems as though self.proc_rank is 0 for the python process on each node at the point that it's attempting to load the spawn weights, so the check introduced to fix #1335 isn't preventing the attempts to load a non existent file.




Discussion:
Hi! thanks for your contribution!, great first issue!
@danpcampbell try 0.7.5?
but yeah this is a bug, it doesn't actually affect anything though... this is just for the end of training that has this annoying crash.

this is done so that .fit() on notebooks restores the new trained weights across processes. However, on a cluster this shouldn't even be called.

The fix is to ONLY call this function if the KAGGLE or COLAB flag are detected.

Mind submitting a PR?
> @danpcampbell try 0.7.5?

Persists in 0.7.5 - I think based on your 2nd comment that this is expected. 

Will next attempt the fix you suggested in 2nd comment, verify, and work up a PR

Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -750,8 +750,9 @@ def fit(
                 # train
                 mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
                 # load weights if not interrupted
-                self.load_spawn_weights(model)
-                self.model = model
+                if os.getenv('COLAB_GPU') or os.getenv('KAGGLE_URL_BASE'):
+                    self.load_spawn_weights(model)
+                    self.model = model
 
         # 1 gpu or dp option triggers training using DP module
         # easier to avoid NCCL issues


Test patch: