Trainer DDP invoking load_spawn_weights() on each node
🐛 Bug

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:

Allocate 4 nodes on a SLURM-managed cluster
srun the script pl.py on each allocated node
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


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.