Shuffling of IterableDataset disabled when dataloader is passed to accelerator.prepare()

Hi, I have an issue with shuffling an IterableDataset that is passed to a dataloader in accelerate. I am trying to use ShufflerIterDataPipe from torch.utils.data.datapipes.iter.combinatorics. As you can see below, it looks like the shuffling is disabled:

from accelerate import Accelerator
from datasets import load_dataset
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe

data = load_dataset("lvwerra/codeparrot-clean-valid", streaming=True, split="train")
shuffled_data = ShufflerIterDataPipe(data, buffer_size=100)
dataloader = DataLoader(shuffled_data, batch_size=4, shuffle=True)
print(f"Before accelerate prepare: {dataloader.dataset._shuffle_enabled}")

accelerator = Accelerator()
dataloader = accelerator.prepare(dataloader)
print(f"After accelerate prepare: {dataloader.dataset._shuffle_enabled}")

Output:

Before accelerate prepare: True
After accelerate prepare: False

I think it’s because ShufflerIterDataPipe returns a new Pytorch IterableDataset that is also an IterDataPipe. When such instance is passed to a Pytorch dataloader, for the shuffling to be enabled we must specify shuffle=True in the dataloader, the distinction is done here. I don’t think this is handled during the dataloader preparation in accelerate.

Python 3.9.12
torch=1.11.0
accelerate=0.8.0

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MULTI_GPU
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

[sgugger]

Happy to review!

[loubnabnl]

It works with torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle) I can make a PR to add this