Lightning+DDP: use_distributed_sampler=True always shuffles data in DDP despite using custom sampler

[dmholtz]

Bug description

When using PyTorch Lightning Trainer in DDP with use_distributed_sampler=True and providing a custom sampler to the dataloader, the sampling order from the custom sampler is neglected. Instead, data is always shuffled.

One can track this issue down to this line where the custom sampler is wrapped with a DistributedSamplerWrapper. However, the kwargs contain shuffle=True due to this line.

When using a custom trainer in combination with fabric as suggested here, the sampling order from the custom sampler is respected. The reason for this is found here where fabric does not pass shuffle=True as keyword argument.

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

from torch.utils.data import Dataset, Sampler, DataLoader
import lightning as L
import torch.nn as nn
import torch

class IntegerDataset(Dataset):
    def __init__(self):
        self.data = [i for i in range(100)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return {"data": torch.tensor([self.data[idx]])}

class InOrderSampler(Sampler):
    def __init__(self, dataset):
        self.dataset = dataset

    def __iter__(self):   
       yield from range(len(self.dataset))

    def __len__(self):
        return len(self.dataset)
    
dataset = IntegerDataset()
sampler = InOrderSampler(dataset)
dataloader = DataLoader(dataset=dataset, batch_size=3, sampler=sampler)

class MyModule(L.LightningModule):

    def __init__(self):
        super().__init__()

        self.layer = nn.Linear(10, 10)

    def training_step(self, batch, batch_idx):
        print(batch)

        input = torch.randn(10, 10, device = self.layer.weight.device)
        output = self.layer(input)
        loss = nn.functional.mse_loss(output, input)
        return loss

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

model = MyModule()

trainer = L.Trainer(
    max_epochs=1,
    use_distributed_sampler=True,
)
trainer.fit(
    model=model,
    train_dataloaders=dataloader,
)

Error messages and logs

The integers from the dataset are not printed in order, but are randomly shuffled.

Environment

Current environment

pytorch-lightning         2.5.1.post0
torch                     2.7.1
torchmetrics              1.7.0
torchvision               0.22.1

More info

No response

cc @justusschock @tchaton

[bhimrazy]

I’d like to dig deeper into this issue. Assigning myself.

[SkafteNicki]

@bhimrazy just wanted to mention that this seems to be a duplicate of #20253

[deependujha]

Hi @dmholtz, thanks for opening the issue.

From the Trainer’s docstring on use_distributed_sampler:

use_distributed_sampler: Whether to wrap the DataLoader's sampler with
                `torch.utils.data.DistributedSampler`.

                If not specified this is toggled automatically for strategies that require it.
                 
                By default, it will add ``shuffle=True`` for the train sampler and
                ``shuffle=False`` for validation/test/predict samplers.

                If you want to disable this logic, you can pass
                ``False`` and add your own distributed sampler in the dataloader hooks. If ``True`` and a distributed
                sampler was already added, Lightning will not replace the existing one...

So in your case you’ll want to:

1. Set use_distributed_sampler=False in the Trainer.
2. Provide your own sampler implementation that handles the distributed logic.

Here’s a minimal working example you can try out:

from torch.utils.data import Dataset, Sampler, DataLoader
import lightning as L
import torch.nn as nn
import torch

class IntegerDataset(Dataset):
    def __init__(self):
        self.data = [i for i in range(100)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        # print(f"Getting item {idx}")
        return {"data": torch.tensor([self.data[idx]])}

class InOrderDistributedSampler(Sampler):
    def __init__(self, dataset, trainer, shuffle=False):
        self.dataset = dataset
        self.global_rank = trainer.global_rank
        self.world_size = trainer.world_size
        self.shuffle = shuffle

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        if self.shuffle:
            # Optionally shuffle the indices
            g = torch.Generator()
            g.manual_seed(0)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        
        # yield from indices: starting from self.global_rank, step by self.world_size
        return iter(indices[self.global_rank::self.world_size])

    def __len__(self):
        return len(self.dataset)



    

class MyModule(L.LightningModule):

    def __init__(self):
        super().__init__()

        self.layer = nn.Linear(10, 10)

    def training_step(self, batch, batch_idx):
        print(f"{self.global_rank=} {batch['data']=}")

        input = torch.randn(10, 10, device = self.layer.weight.device)
        output = self.layer(input)
        loss = nn.functional.mse_loss(output, input)
        return loss

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


if __name__ == "__main__":
    trainer = L.Trainer(
        max_epochs=1,
        use_distributed_sampler=False,
    )

    dataset = IntegerDataset()
    sampler = InOrderDistributedSampler(dataset, trainer)
    dataloader = DataLoader(dataset=dataset, batch_size=3, sampler=sampler)

    model = MyModule()

    trainer.fit(
        model=model,
        train_dataloaders=dataloader,
    )

This way, Lightning won’t inject its own DistributedSampler, and your sampler fully controls how data is partitioned across ranks.

Thanks.

cc: @bhimrazy

[dmholtz]

@deependujha Thank you very much for this detailed response, this helps.

Initially, it seems a bit counter-intuitive to obtain rank and world_size from the trainer when building the sampler, but it works.

Since the shuffling behavior is apparently consistent with the docstring, I think this issue can be closed.