batch sampler, ddp, and sampler wrapper

[Jserax]

Bug description

I might be wrong, but when using a BatchSampler with DDP, Lightning never wraps it with a DistributedSamplerWrapper. This is due to the following condition in its code:

if isinstance(dataloader.sampler, (RandomSampler, SequentialSampler)):
    return DistributedSampler(dataloader.dataset, **kwargs)

Since PyTorch uses a SequentialSampler by default, this condition is always met

DataLoader(ds, sampler=None, batch_sampler=MNSampler(ds.labels)).sampler
<torch.utils.data.sampler.SequentialSampler at 0x7deb78e01eb0>

What version are you seeing the problem on?

v2.4

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

cc @justusschock @lantiga

[bhimrazy]

Hi @Jserax, Thanks for reporting!
Could you please share a minimal reproducible example that shows the issue? That would help us confirm whether the sampler isn’t being wrapped correctly in DDP and speed up debugging.

[SkafteNicki]

If you are providing a BatchSampler or a subclass of BatchSampler then lightning will still wrap it, it happens in just the line beneath the lines you linked:

pytorch-lightning/src/lightning/pytorch/trainer/connectors/data_connector.py

Lines 241 to 243 in e55650d

	if isinstance(dataloader.sampler, (RandomSampler, SequentialSampler)):
	return DistributedSampler(dataloader.dataset, **kwargs)
	return DistributedSamplerWrapper(dataloader.sampler, **kwargs)

e.g. RandomSampler and SequentialSampler gets wrapped in DistributedSampler and dataloaders not fulfilling these criteria are then wrapped using DistributedSamplerWrapper

[SkafteNicki]

Closing issue, but feel free to ping and reopen if more information is provided.