Replace custom samplers with distributed versions when replace_sampler_ddp=True
🚀 Feature
Replace custom samplers with distributed versions when replace_sampler_ddp=True.
Motivation
One of the issues I've been having is using custom samplers with distributed training. Prior to Lightning, I was using the DistributedSamplerWrapper class from Catalyst (https://catalyst-team.github.io/catalyst/_modules/catalyst/data/sampler.html) as part of my own pipeline to use custom samplers during distributed training.
When I try to define my own custom DistributedSampler when using Lightning, I always get the following error:
AssertionError: Default process group is not initialized
It makes sense because the data loader gets called before distributed is initialized. (See: #3238 (comment))
However unlike in the comment, I'm unable to create the DistributedSampler by just leaving the arguments as their default since I always get the above AssertionError. It looks like the original poster was using TPUs, and it seems that on TPUs you can pass rank and num_replicas to the sampler without having to run torch.distributed.init_process_group().
Pitch
Basically, what I do now is I manually edit the TrainerDataLoadingMixin class in trainer/data_loading.py.
I change the auto_add_sampler() method:
if self.replace_sampler_ddp and need_dist_sampler:
    if not isinstance(dataloader.sampler, (SequentialSampler, RandomSampler)):
        rank_zero_warn('Replacing custom sampler with distributed version ...')
        sampler = self._wrap_distributed_sampler(dataloader, shuffle)
    else:
        sampler = self._get_distributed_sampler(dataloader, shuffle)
dataloader = self.replace_sampler(dataloader, sampler)

New _wrap_distributed_sampler() method:
def _wrap_distributed_sampler(self, dataloader, shuffle):
    kwargs = self.distributed_sampler_kwargs
    kwargs['shuffle'] = shuffle and not self.overfit_batches
    sampler = DistributedSamplerWrapper(dataloader.sampler, **kwargs)
    return sampler

Alternatives
There may be a simple workaround that I'm missing for why I'm not able to pass in my own custom distributed samplers, in which case this wouldn't be needed.