replace_sampler_ddp + batch_sampler

[sshleifer]

Is it possible to make a distributed-friendly batch_sampler that gets passed to DataLoader.

I have a working single-gpu version that produces an iterator where each .next is list of ids, but when I pass it with

            return DataLoader(
                dataset,
                batch_sampler=batch_sampler,
                collate_fn=dataset.collate_fn,

                num_workers=self.num_workers,
             
            )

I get too few batches and each batch is way to small.

• Without --gpus 2 --replace_sampler_ddp everything is OK.
• With --gpus 2 alone I get the expected error.

Has anyone seen a working example of this?

[ananyahjha93]

@sshleifer
I would probably use https://pytorch.org/docs/stable/_modules/torch/utils/data/distributed.html#DistributedSampler as a reference. Have you included this step?

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples