DDP on iterable Dataset?

Ravi_Teja2 · July 13, 2024, 8:19am

I’m bit new to using Iterable Datasets. I have been using torchdata to build my dataloaders but they seem to deprecate and even delete this functionality now . So, I want to build Iterable Dataset but it seems DistributedSampler cannot be used on iterable dataset. Any suggestions on how to use DDP on iterable Datasets?

I’m aware of this large issue: ChunkDataset API proposal by thiagocrepaldi · Pull Request #26547 · pytorch/pytorch · GitHub
But I don’t think this functionality is added yet.
And also this post: Using IterableDataset with DistributedDataParallel
Here I don’t find any workable process or examples

602389789 · January 8, 2025, 2:10am

Maybe you can use something like the IterableDatasetShard to wrap your IterableDataset

github.com

huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L798


      
                  for batch_start in range(self.batch_size * self.process_index, self.total_num_samples, self.total_batch_size):
                      result += indices[batch_start : batch_start + self.batch_size]
          
                  return iter(result)
          
              def __len__(self):
                  # Each shard only sees a fraction of total_num_samples.
                  return self.total_num_samples // self.num_processes
          
          
          class IterableDatasetShard(IterableDataset):
              """
              Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will
              always yield a number of samples that is a round multiple of the actual batch size (which is `batch_size x
              num_processes`). Depending on the value of the `drop_last` attribute, it will either stop the iteration at the
              first batch that would be too small or loop with indices from the beginning.
          
              On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]` with a batch size of
              2:
          
              - the shard on process 0 will yield `[0, 1, 4, 5, 8, 9]` so will see batches `[0, 1]`, `[4, 5]`, `[8, 9]`