Training stalls with DDP and iterable training dataset at validation step for any val_check_interval>1
🐛 Bug
Training stalls with DDP and iterable training dataset at validation step for any val_check_interval>1. It works fine for val_check_interval=1.
To Reproduce
Here is a toy example to reproduce:
import numpy as np
import pybmi.torch as pt
import pytorch_lightning as pl
import torch
from torch import nn
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
from torch.utils.data import IterableDataset, DataLoader, ChainDataset, get_worker_info
import torch.distributed as dist 
    

class MyIterableDataset(IterableDataset):
    
    def __init__(self, ds_list):
        self.seq_len =  3
        self.batch_size = 2
        self.ds_list = ds_list
        
    def process_data(self, worker_ds):
        l = []
        for d in worker_ds:
            l.append(d)
        pdata = np.concatenate(l)
        batch_size_total = self.batch_size * self.seq_len
        n_batches = len(pdata)//batch_size_total
        pdata = pdata[:n_batches * batch_size_total]
        pdata = pdata.reshape((self.batch_size, -1)).astype(np.float32)
        for n in range(0, pdata.shape[1], self.seq_len):
            itr = pdata[:, n:n+self.seq_len] 
            yield itr
    
    def __iter__(self):
        rank = dist.get_rank()
        world_size = dist.get_world_size()
        worker_info = get_worker_info()
        total_chunks = world_size * worker_info.num_workers
        chunk_id = (rank * worker_info.num_workers) + worker_info.id        
        ds_list = np.array(self.ds_list,dtype = 'object')
        split = np.array_split(ds_list, total_chunks)
        return (self.process_data(split[chunk_id]))
    
class TestDataset():
    def __init__(self):   
        self.maxlen = 40
    def process_data(self,ds):
        out = np.zeros([self.maxlen,len(ds)])
        n = []
        for ii,pdata in enumerate(ds):
            n.append(len(pdata))
            out[:n[-1],ii] = pdata
        return out[:max(n),:].astype(np.float32)
    
a = np.arange(20)*1.
b = np.arange(23)*-1.
c = np.arange(31)*0.01
d = np.arange(17)*-.01
data = [a,b,c,d]


class PnRDataModule(pl.LightningDataModule):
    
    def __init__(self):
        super().__init__()
        self.batch_size = None
        self.val_batch_size = 3
        self.num_workers = 2
        
    def setup(self, stage=None):
        if stage == 'fit' or stage is None:
            self.training_data = MyIterableDataset(data)
            self.val_data = TestDataset().process_data(data)
          
    def train_dataloader(self):
        return DataLoader(
            self.training_data,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )
    def val_dataloader(self): 
        return DataLoader(
            self.val_data,
            batch_size=self.val_batch_size,
            num_workers=2,
            pin_memory=True,
            drop_last = True,
        )

class Model(nn.Module):
    
    def __init__(self):
        super(Model, self).__init__()
        self.ff1 = nn.Linear(1,100)
        self.ff2 = nn.Linear(100,1)
    def forward(self, x):
        x1 = self.ff1(x)
        y = self.ff2(x1)
        return y
    
    
class LM(pl.LightningModule):
    
    def __init__(self):
        super().__init__()
        self.model = Model()
        self.criterion = torch.nn.MSELoss(reduction='mean')
    def forward(self,x):
        logits= self.model.forward(x)
        return logits
    
    def training_step(self, batch, batch_idx):
        batch = batch.unsqueeze(2)
        batch = batch.permute([1,0,2])
        logits = self.model.forward(batch)
        loss = self.criterion(logits, batch) 
        self.log('train_loss', loss, prog_bar=True)
        return  loss

    def validation_step(self, batch, batch_idx):
        logits = self.model.forward(batch.unsqueeze(2))
        loss = self.criterion(logits.squeeze(2), batch) 
        self.log('val_loss', loss, prog_bar=True)
        return loss
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        return optimizer

def main():

    dm = PnRDataModule()
    model = LM()
    checkpoint_callback = ModelCheckpoint(dirpath='Models/MMD',monitor="val_loss", mode="min")
    trainer = pl.Trainer(max_epochs=10, gpus=2,callbacks=[checkpoint_callback], accelerator='ddp',val_check_interval=2)
    trainer.fit(model, dm)

if __name__ == "__main__":
    main()

Expected behavior
Training to work with any val_check_interval.
Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !

PyTorch Version (e.g., 1.0): 1.8.0
OS (e.g., Linux): Linux
How you installed PyTorch (conda, pip, source): conda
Build command you used (if compiling from source):
Python version: 3.6.10
CUDA/cuDNN version: 11.1
GPU models and configuration: 4x V100 GPU
Any other relevant information:

Additional context

cc @tchaton