Interesting thread.
If we want to do a simple random sampler, wouldn’t something like this work? We just take a random sample of our whole dataset on point of dataset creation, than wrap our Dataset class in DistributedSampler, which would take care of splitting it among processes?
We make reload  of the Dataset every epoch so a new sample is drawn.
class Dataset():
    def __init__(self, seed, sample_size, *args, **kwargs):
        random.seed(seed)
        random.shuffle(self.ids)
        self.ids = self.ids[:sample_size]

dataset = Dataset()
dataset = torch.utils.data.distributed.DistributedSampler(dataset)