DistributedDataParallel with nccl backend produces zombie processes
Hey lightning community,
first I want to thank you for this nice project. It helped me a lot to improve my research code and I'm happy to recommend it to my colleagues whenever they are complaining about their code mess.
I have a problem with some kind of racing condition which is reproducible in a slurm environment with multiple GPUs. I used DistributedDataParallel with the 'nccl'-backend.
The default implementation of PyTorch-lightning can produce zombie processes, which reserve GPU memory and prevent further usage of the memory. It happens mainly if the main process is stopped or crashes.
You can reproduce this with the given code. If the training starts and the main process is killed with a strict signal like SIGKILL, the child processes stay persistent in most cases.
I could solve the problem for me by overwriting init_ddp_connection. The overwritten method sets the NCCL_BLOCKING_WAIT=1 and reduced the timeout additionally. (The torch documentation mentions that the timeout is only used for nccl if NCCL_BLOCKING_WAIT is 1)
Is there a better way to get rid of this problem?
Or should we adjust the default behavior of lightning for the nccl backend?
Best regrads
Leon Varga
Code
import pytorch_lightning as lightning
import torch
from torch.utils.data import DataLoader


class MyModule(lightning.LightningModule):
    def __init__(self):
        super(MyModule, self).__init__()

        self.model = torch.nn.Linear(1000, 1000)
        self.criterion = torch.nn.MSELoss()

    def train_dataloader(self) -> DataLoader:
        data = torch.randn((int(1e5), 2, 1000))

        training_generator = DataLoader(data)
        return training_generator

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters())

        return [optimizer]

    def training_step(self, batch, batch_idx):
        pred = self.model(batch[:, 0])
        loss = self.criterion(pred, batch[:, 1])

        return {
            'loss': loss,
        }

    def forward(self, x):
        return self.model(x)



if __name__ == "__main__":
    model = MyModule()
    trainer = lightning.Trainer(default_root_dir='/tmp/test',
                                max_epochs=100, gpus=-1,
                                distributed_backend='ddp')
    trainer.fit(model)

What's your environment?

Linux
Pytorch 1.5 / 1.6
PytorchLightning 0.8.5