Training is interrupted without error with MulitGPU

[skull3r7]

🐛 Bug

The training is interrupted randomly in the middle of an epoch without errors. The console only says: Terminated.
The error does not necessarily occur, if it does then mostly between epochs 2-4. It is noticeable that processes are still running after the termination, the graphic cards are still used by python processes.

We train the PyTorch version of the ImageGPT model with huggingface transformers. Could also be problem of huggingface, we are not sure.

Epoch 1: 29%|█▍ | 9413/32393 [3:28:18<8:28:33, 1.33s/it, loss=3.23, v_num=9]Terminated

Please reproduce using the BoringModel

Cant reproduce with Boring Model.

Code

class ImageGPT(pl.LightningModule):

    def __init__(self,
                 learning_rate=learning_rate
                 ):
        super().__init__()
        self.gpt2 =  ImageGPT2LMHeadModel(config=...)
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.learning_rate = learning_rate

    def forward(self, x):
        return self.gpt2(x, past_key_values=None)

....


logger = pl_loggers.TensorBoardLogger(save_dir="logs", name=name)

checkpoint_callback = ModelCheckpoint(
        save_top_k=1,
        verbose=True,
        monitor='val_loss',
        mode='min',
        filepath='../models',
        prefix='ImageGPT'
    )

trainer = Trainer(
                accelerator='ddp',
                max_epochs=10,
                max_steps=None,
                precision=32,
                accumulate_grad_batches=1,
                gpus=[0, 1, 2],
                callbacks=[checkpoint_callback],
                logger=logger,
                gradient_clip_val=0.6
            )

trainer.fit(model=model, datamodule=datamodule)

Expected behavior

The training is fully completed across all epochs.

Environment

• CUDA:
  • GPU:
    • TITAN RTX
    • TITAN RTX
    • TITAN RTX
  • available: True
  • version: 10.2
• Packages:
  • numpy: 1.19.4
  • pyTorch_debug: False
  • pyTorch_version: 1.7.1
  • pytorch-lightning: 1.1.2
  • transformers: 3.5.1
  • tqdm: 4.55.0
• System:
  • OS: Linux, 64bit
  • processor: x86_64
  • python: 3.7.4
  • version: 86-Ubuntu SMP Fri Jan 17 17:24:28 UTC 2020

Additional context

We have made the following points to solve the problem:

• set the num-workers of the dataloaders to 0 or 1 (instead of 32-64)
• go back to 32 bit precision
• different learning rates
• added gradient clipping
• used AdamW implementation from huggingface

[angadkalra]

This happens to me too, but I don't get "Terminated" at end of progress bar, it just stops, and when I check the system with "top i", I see 4 python processes running at 100% and 4 GPUs at about 90% capacity but nothing is changing. Sometimes, after like an hour or two, it just randomly restarts but the s/it jumps from 2.5 to like 85.

[XiaomoWu]

@angadkalra +1. For me switching to DP solves the problem, but it at the cost of speed.

Update:
I test on PL v1.1.5-v1.2.0 rc1, and the problem persists. I'm sorry I can't upload a reproducible example at the moment, but will probably do that later.

[carmocca]

Hi everybody! If any of you can provide a snippet that would be awesome. Otherwise we are blind trying to fix it.

[angadkalra]

Hi everybody! If any of you can provide a snippet that would be awesome. Otherwise we are blind trying to fix it.

Code snippet? Is the code in the original post good enough or need colab?

[carmocca]

Need something I can run. Preferably a colab 👍

[XiaomoWu]

@carmocca Sorry for the late reproducible example! Please see below for a self-contained example that uses two GPUs for DDP. For me the code gets stuck at epoch 13, while the two GPUs keep busy at 100%. Switching to DP solves the problem.

Some hints

• Removing the dropout layer in the def shared_step(self, batch) method also solves the problem. But how could such a harmless dropout layer ruin the model?
• I can reproduce the bug on both lightning 1.1.8 and 1.2.0

My settings:

• OS: Ubuntu 20.10
• pytorch: 1.7.1 + CUDA 11.0
• pytorch-lightning == 1.1.8
• GPUs: 2080 Ti x2

import torch
import pytorch_lightning as pl
import torch.nn.functional as F
import torch.optim as optim

from torch import nn
from torch.utils.data import Dataset, DataLoader

# set random seed
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True

# Define Dataset
class CCDataset(Dataset):
    
    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return self.data.shape[0]
    
    def __getitem__(self, idx):
        target = self.data[idx,0].clone().detach()
        features = self.data[idx,1:].clone().detach()

        return target, features


# then define DataModule
class CCDataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        
    # Dataset
    def setup(self):
        # read the train and test dataset
        # targets_train = feather.read_feather('targets_train.feather')
        # targets_val = feather.read_feather('targets_test.feather')
        targets_train = torch.rand(12157,16)
        targets_val = torch.rand(12157,16)
        
        self.train_dataset = CCDataset(targets_train)
        self.val_dataset = CCDataset(targets_val)

    # DataLoader
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=64, 
                          shuffle=True, drop_last=False, num_workers=2,
                          pin_memory=True)
    
    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=64,
                          num_workers=2, pin_memory=True,
                          drop_last=False)


# # def Model

class Model(pl.LightningModule):
    '''Mainly define the `*_step_end` methods
    '''
    def __init__(self):
        super().__init__()
        
        # dropout layers
        self.dropout_1 = nn.Dropout(0.5)
        
        # fc layers
        self.fc_1 = nn.Linear(15, 16)
        self.fc_2 = nn.Linear(16, 1)
        
    def shared_step(self, batch):
        t, x = batch
        x = self.dropout_1(F.relu(self.fc_1(x)))
        y = self.fc_2(x) # (N, 1)    
        
        return y.squeeze(), t
        
    # train step
    def training_step(self, batch, idx):
        y, t = self.shared_step(batch)
        return {'y': y, 't': t}
        
    # validation step
    def validation_step(self, batch, idx):
        y, t = self.shared_step(batch)
        return {'y': y, 't': t}
        
    # loss
    def mse_loss(self, y, t):
        return F.mse_loss(y, t)
        
    # def training_step_end
    def training_step_end(self, outputs):
        y = outputs['y']
        t = outputs['t']
        loss = self.mse_loss(y, t)
        
        return {'loss':loss}
    
    # def validation_step_end
    def validation_step_end(self, outputs):
        y = outputs['y']
        t = outputs['t']
        
        return {'y': y, 't': t}
        
    # validation step
    def validation_epoch_end(self, outputs):
        y = torch.cat([x['y'] for x in outputs])
        t = torch.cat([x['t'] for x in outputs])
        
        loss = self.mse_loss(y, t)
        rmse = torch.sqrt(loss)
        self.log('val_rmse', rmse, on_step=False)
        
    # optimizer
    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=1e-4)
        return optimizer  


# # Run

# checkpoint
checkpoint_callback = pl.callbacks.ModelCheckpoint(
    verbose=True,
    mode='min',
    monitor='val_rmse',
    save_top_k=1)

# trainer
trainer = pl.Trainer(gpus=[0,1], 
                     checkpoint_callback=checkpoint_callback, 
                     accelerator='ddp',
                     min_epochs=10,
                     max_epochs=500)

# loop over windows
torch.manual_seed(42)

# init model
model = Model()

# create datamodule
datamodule = CCDataModule()
datamodule.setup()

# train the model
trainer.fit(model, datamodule)