Validation loader is seeded the same during epoch in distributed mode

Bug description

The validation dataloader is seeded identically for each validation step during each epoch in distributed mode. While this may be fine when the whole validation set is consumed entirely in the validation step, this leads to unexpected results when using the limit_val_batches option.

It seems to me that this is not the result that is expected when setting limit_val_batches ; it seems that that should give a different subset of batches every time the validation loop is executed. The behavior for devices=1 confirms this: in this case the dataloader provides a new set of batches on every validation loop, even in the case of multiple validation loops per epoch.

The issue seems to arise from the distributed sampler in pytorch (torch.utils.data.distributed.DistributedSampler, see https://pytorch.org/docs/stable/_modules/torch/utils/data/distributed.html#DistributedSampler), which seeds based on a seed, and epoch number. The effect is that the exact same shuffling is used as long as the epoch number has not been updated. And then, as mentioned before, when limiting the validation to a small subset of the total validation set, over the course of one epoch, exactly the same batches are used for evaluation at every step, and then a new fixed set on the second epoch, etcetera.

Below is a reproducing example that displays the batch indices at every validation step. It is easy to see that the output is identical all throughout the first epoch. Setting the number of devices to 1, gives a random set of batches every validation step (the expected behavior).

Summary:
The combination devices > 1, limit_val_batches < (len(dataset) // batch_size) and shuffle=True gives the same set of validation batches for each validation step.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

from torch.utils.data import DataLoader
import torch
import torch.nn as nn
import torch.optim as optim
import pytorch_lightning as pl
import time

# Define the LightningModule
class MyModel(pl.LightningModule):
    def __init__(self):
        super(MyModel, self).__init__()
        self.model = nn.parameter.Parameter(torch.ones(1,), requires_grad=True)

    # Nothing important here
    def training_step(self, batch, batch_idx):
        loss = (self.model - batch).square().mean()
        self.log('train_loss', loss)
        return loss

    def train_dataloader(self):
        dataset = torch.range(0, 100000).float() / 10000
        train_dataloader = DataLoader(
            dataset, batch_size=2)
        return train_dataloader

    def val_dataloader(self):
        val_dataset = torch.range(0, 10000).float()
        val_batch_size=8

        val_dataloader = DataLoader(
            val_dataset, batch_size=val_batch_size, drop_last=False, num_workers=0,
            pin_memory=False, collate_fn=collate_fn, shuffle=True
        )
        return val_dataloader

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=1e-4)
        return optimizer

    def validation_step(self, batch, batch_idx):
        print(f'Reporting from local rank {self.trainer.local_rank}, index {batch_idx}, batch {batch}')

collate_fn = torch.stack

if __name__ == '__main__':
    devices= 2
    accelerator = 'cpu'
    # Create a PyTorch Lightning trainer
    model = MyModel()

    trainer = pl.Trainer(
        max_epochs=100,
        val_check_interval=0.2,
        limit_val_batches=2,
        devices=devices,
        accelerator=accelerator,
    )


    # Train the model using the Lightning trainer
    trainer.fit(model)

Error messages and logs

[ ... ] 
Epoch 0:  20%|███████████████████████▌                                                                                              | 5000/25001 [00:06<00:24, 811.98it/s, v_num=0]
Reporting from local rank 1, index 0, batch tensor([8240., 9125., 2912., 3170., 9948., 9427., 6728., 1879.])
Reporting from local rank 1, index 1, batch tensor([6557., 5502., 8580., 1543., 6386., 1750., 4385., 3793.])
Validation: 0it [00:00, ?it/s]                                                                                                                                                    
Reporting from local rank 0, index 0, batch tensor([ 354., 2137., 5853., 6143., 8100., 3042., 2407., 1548.])                                                  | 0/2 [00:00<?, ?it/s]
                                                                                                                                                                                  
Reporting from local rank 0, index 1, batch tensor([7737., 3251., 6256., 6769., 7153.,  720., 3561., 5881.])                                        | 1/2 [00:00<00:00, 2004.93it/s]
Epoch 0:  40%|██████████████████████████████████████████████▊                                                                      | 10000/25001 [00:12<00:18, 806.40it/s, v_num=0]
Reporting from local rank 1, index 0, batch tensor([8240., 9125., 2912., 3170., 9948., 9427., 6728., 1879.])                                                                       
Reporting from local rank 1, index 1, batch tensor([6557., 5502., 8580., 1543., 6386., 1750., 4385., 3793.])
                                                                                                                                                                                  
Reporting from local rank 0, index 0, batch tensor([ 354., 2137., 5853., 6143., 8100., 3042., 2407., 1548.])                                                  | 0/2 [00:00<?, ?it/s]
                                                                                                                                                                                  
Reporting from local rank 0, index 1, batch tensor([7737., 3251., 6256., 6769., 7153.,  720., 3561., 5881.])                                        | 1/2 [00:00<00:00, 1972.86it/s]
Epoch 0:  60%|██████████████████████████████████████████████████████████████████████▏                                              | 15000/25001 [00:18<00:12, 799.65it/s, v_num=0]
Reporting from local rank 1, index 0, batch tensor([8240., 9125., 2912., 3170., 9948., 9427., 6728., 1879.])                                                                       
Reporting from local rank 1, index 1, batch tensor([6557., 5502., 8580., 1543., 6386., 1750., 4385., 3793.])
                                                                                                                                                                                  
Reporting from local rank 0, index 0, batch tensor([ 354., 2137., 5853., 6143., 8100., 3042., 2407., 1548.])                                                  | 0/2 [00:00<?, ?it/s]
                                                                                                                                                                                  
Reporting from local rank 0, index 1, batch tensor([7737., 3251., 6256., 6769., 7153.,  720., 3561., 5881.])                                        | 1/2 [00:00<00:00, 1897.88it/s]
[ ... ] 

Environment

More info

No response

cc @justusschock @awaelchli