Prefetching for IterableDataset
State: open
Labels: description: New feature or request
name: enhancement

### Feature request

Add support for prefetching the next n batches through iterabledataset to reduce batch loading bottleneck in training loop.

### Motivation

The primary motivation behind this is to use hardware accelerators alongside a streaming dataset. This is required when you are in a low ram or low disk space setting as well as quick iteration where you're iterating though different accelerator environments (e.x changing ec2 instances quickly to figure out batch/sec for a particular architecture). 

Currently, using the IterableDataset results in accelerators becoming basically useless due to the massive bottleneck induced by the dataset lazy loading/transform/mapping.

I've considered two alternatives:
PyTorch dataloader that handles this. However, I'm using jax, and I believe this is a piece of functionality that should live in the stream class.

Replicating the "num_workers" part of the PyTorch DataLoader to eagerly load batches and apply the transform so Arrow caching will automatically cache results and make them accessible.

### Your contribution

I may or may not have time to do this. Currently, I've written the basic multiprocessor approach to handle the eager DataLoader for my own use case with code that's not integrated to datasets. I'd definitely see this as being the default over the regular Dataset for most people given that they wouldn't have to wait on the datasets while also not worrying about performance.

Comments:
Very cool! Do you have a link to the code that you're using to eagerly fetch the data? Would also be interested in hacking around something here for pre-fetching iterable datasets

I ended up just switching back to the pytorch dataloader and using it's multiprocessing functionality to handle this :(. I'm just not that familiar with python multiprocessing to get something to work in jupyter (kept having weird behaviors happening with zombies living after the cell finished).

Ultimately settled on using webdataset to circumvent huggingface datasets entirely. Would definitely switch back if: https://github.com/huggingface/datasets/issues/5337 was resolved.

Hi! You can combine `datasets` with `torchdata` to prefetch `IterableDataset`'s samples:
```python
from datasets import load_dataset
from torchdata.datapipes.iter import IterableWrapper, HuggingFaceHubReader
from torch.utils.data import DataLoader

ds = load_dataset("sst", split="train", streaming=True)
# processing...
dp = IterableWrapper(ds)
dp = dp.prefetch(100)
dl = DataLoader(dp, batch_size=8)

i = iter(dl)
next(i)
```

Hey @mariosasko! Thanks for the tip here - introducing prefetch with `torchdata` didn't really give me any performance difference vs not prefetching, but the concept is definitely one that could be really beneficial. Are there any benchmarks that show the speed-up you can get with `torchdata`'s prefetch just for comparison?

I want to perform mapping() in advance by prefetching within the IterableDataset. Are there any recent updates?

> Ultimately settled on using webdataset to circumvent huggingface datasets entirely. Would definitely switch back if: [#5337](https://github.com/huggingface/datasets/issues/5337) was resolved.

The HF dataset now supports the webdataset format. On a custom dataset with multiple .tar files that is on a single HDD. The read speed is ~180 MB/s when building the buffer.

Using the following setup, I observe ~170MB/s read speed when training with hf accelerate (multi-gpu)

```
    train_dataset = load_dataset(
        "path_to_py",
        split="train",
        streaming=True,
        trust_remote_code=True
    )
    train_dataset = train_dataset.shuffle(seed=general_config.seed, buffer_size=1000)
    train_dataset = train_dataset.with_format("torch")
    train_dataloader = DataLoader(train_dataset, num_workers=1, batch_size=8, prefetch_factor=4)
   # Then accelerator.prepare (I used dispatch_batches=False as my custom dataset does not return tuple of tensors)
```


> > Ultimately settled on using webdataset to circumvent huggingface datasets entirely. Would definitely switch back if: [#5337](https://github.com/huggingface/datasets/issues/5337) was resolved.
> 
> The HF dataset now supports the webdataset format. On a custom dataset with multiple .tar files that is on a single HDD. The read speed is ~180 MB/s when building the buffer.
> 
> Using the following setup, I observe ~170MB/s read speed when training with hf accelerate (multi-gpu)
> 
> ```
>     train_dataset = load_dataset(
>         "path_to_py",
>         split="train",
>         streaming=True,
>         trust_remote_code=True
>     )
>     train_dataset = train_dataset.shuffle(seed=general_config.seed, buffer_size=1000)
>     train_dataset = train_dataset.with_format("torch")
>     train_dataloader = DataLoader(train_dataset, num_workers=1, batch_size=8, prefetch_factor=4)
>    # Then accelerator.prepare (I used dispatch_batches=False as my custom dataset does not return tuple of tensors)
> ```

Thanks for the tip!

I am wondering. How large the shard files should be i.e. the tar files are?
What is the memory peak used -- the size of the single shard file? 

Do you know if HF load_dataset uses the same shard if used with DDP or more of them as a) DDP number of GPU processes increase b) the number of workers within the single ddp_rank process increase?

Could you please point me to minimalistic example of using webdatasets which you worked with?
I looked at few posts at HF forum, but most posts raised more questions than provided answers.

Hi,

> I am wondering. How large the shard files should be i.e. the tar files are? What is the memory peak used -- the size of the single shard file?
I forgot the exact recommendation, you shouldn't worry too much about the number of tar files, just making sure the files in tar are chunked reasonably to avoid overhead (don't dump a lot of small files in tar) and you should easily reach I/O bottleneck. I tested my implementation on HPC and got 1GB/s read, seems limiting by 10G Ethernet.

In terms of the memory, the formula that I observed: `Total Mem formula: num_worker*num_GPU*Buffer`, num_worker refers to the dataloader, buffer example: A single snapshot is ~9.8MB, the pre-allocated buffer contains 1000 snapshots (for shuffling) -> ~10G. The buffer size affect the randomness

> 
> Do you know if HF load_dataset uses the same shard if used with DDP or more of them as a) DDP number of GPU processes increase b) the number of workers within the single ddp_rank process increase?

I think the formula above is what you asked

> Could you please point me to minimalistic example of using webdatasets which you worked with? I looked at few posts at HF forum, but most posts raised more questions than provided answers.

You can take a look at my [repo](https://github.com/tonyzyl/ladcast/blob/master/ladcast/dataloader/weather_dataset.py). Not the minimal example though, I follow the official legacy webdataset example