Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#21,131 16 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug data handling distributed ver: 2.5.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### 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](https://github.com/Lightning-AI/pytorch-lightning/blob/634e6e6d06fced4cacfb9de675fc328ace1a06c4/src/lightning/pytorch/trainer/connectors/data_connector.py#L243) where the custom sampler is wrapped with a `DistributedSamplerWrapper`. However, the kwargs contain `shuffle=True` due to [this line](https://github.com/Lightning-AI/pytorch-lightning/blob/e55650dcebe1cb9c50addab06d292a1704b1d288/src/lightning/pytorch/trainer/connectors/data_connector.py#L490).

When using a custom trainer in combination with fabric as suggested [here](https://github.com/Lightning-AI/pytorch-lightning/tree/master/examples/fabric/build_your_own_trainer), the sampling order from the custom sampler is respected. The reason for this is found [here](https://github.com/Lightning-AI/pytorch-lightning/blob/e55650dcebe1cb9c50addab06d292a1704b1d288/src/lightning/fabric/fabric.py#L398) 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

```python
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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/lightning/pytorch/trainer/connectors/data_connector.py at the linked sampler-wrapping and kwargs lines, then compare the corresponding logic in src/lightning/fabric/fabric.py. Reproduce the behavior with the provided custom sampler and DDP example; done means the custom sampling order is preserved when use_distributed_sampler=True.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.