Problem in multi-node training
Hello pytorch-lightning community,
my training hangs when training on multi-nodes; on single node with multiple GPUs runs fine :/
It baffles me that although the global rank ID seems right, the member output has 4 instead of 8 in the denominator.
Since I run in a slurm environment, do I have to add the SLURMEnvironment plugin in the Trainer? I tried to add it alongside the DDPPlugin but it was not accepted (Found invalid type for plugin <class 'pytorch_lightning.plugins.environments.slurm_environment.SLURMEnvironment'>. Expected a precision or training type plugin)
The job submission file has the corresponding lines:
#SBATCH --gres=gpu:4
#SBATCH --nodes=2
#SBATCH --exclusive
srun --ntasks=8 python3 coolModel.py 2>&1 | tee log.train
I attach the output and the code below...
pytorch version:  1.8.1
pytorch-lightning version:  1.2.4
Cheers,
Nikos
multNodeTraining.txt
###########python code
import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
import pytorch_lightning as pl
from pytorch_lightning import Trainer
#from test_tube import Experiment
class database(pl.LightningDataModule):
def __init__(self):
    super().__init__()

#def setup(self, stage=None):
def train_dataloader(self):
    return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

def val_dataloader(self):
    return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

def test_dataloader(self):
    return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

class CoolModel(pl.LightningModule):
def __init__(self):
    super(CoolModel, self).__init__()
    # not the best model...
    self.l1 = torch.nn.Linear(28 * 28, 10)

def forward(self, x):
    return torch.relu(self.l1(x.view(x.size(0), -1)))

def my_loss(self, y_hat, y):
    return F.cross_entropy(y_hat, y)

def training_step(self, batch, batch_nb):
    x, y = batch
    y_hat = self.forward(x)
    return {'loss': self.my_loss(y_hat, y)}

def validation_step(self, batch, batch_nb):
    x, y = batch
    y_hat = self.forward(x)
    return {'val_loss': self.my_loss(y_hat, y)}

def validation_end(self, outputs):
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
    return {'avg_val_loss': avg_loss}

def configure_optimizers(self):
    return [torch.optim.Adam(self.parameters(), lr=0.02)]

if name=='main':
dm = database()
model = CoolModel()
#exp = Experiment(save_dir=os.getcwd())
#checkp1 = pl.callbacks.ModelCheckpoint(
#    monitor='loss', save_top_k=2, dirpath='./', mode='min', save_last=True
#    ) 
trainer = Trainer(
        max_epochs=10,
        gpus=4, num_nodes=2, accelerator='ddp',
        plugins= pl.plugins.DDPPlugin(find_unused_parameters=False), #pl.plugins.SLURMEnvironment],
        #progress_bar_refresh_rate=0
        )
# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)
trainer.fit(model, dm)
# view tensorflow logs
print(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
print('and going to http://localhost:6006 on your browser')
Hello Nikos
Do you have 8 gpus in the node? I think it must match gres.
Don't you also need to specify how many tasks per node in the SBATCH directive? [1]
Also, I notice some unsupported Trainer arguments in your script. It should be:
trainer = Trainer(max_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], num_nodes=4) 
Make sure this script actually runs on CPU first before going to the cluster 😅
Totally no slurm expert here, just looking at your script with one eye closed.
[1] https://pytorch-lightning.readthedocs.io/en/latest/clouds/cluster.html#slurm-managed-cluster