Question about iterable dataset and streaming
Author: Tavish9
State: open

In the doc, I found the following example: https://github.com/huggingface/datasets/blob/611f5a592359ebac6f858f515c776aa7d99838b2/docs/source/stream.mdx?plain=1#L65-L78

I am confused, 
1. If we have already loaded the dataset, why doing `to_iterable_dataset`?  Does it go through the dataset faster than map-style dataset?
2. `load_dataset(streaming=True)` is useful for huge dataset, but the speed is slow. How to make it comparable to `to_iterable_dataset` without loading the whole dataset into RAM?

Comments: > If we have already loaded the dataset, why doing to_iterable_dataset? Does it go through the dataset faster than map-style dataset?

yes, it makes a faster DataLoader for example (otherwise DataLoader uses `__getitem__` which is slower than iterating)

> load_dataset(streaming=True) is useful for huge dataset, but the speed is slow. How to make it comparable to to_iterable_dataset without loading the whole dataset into RAM?

You can aim for saturating your bandwidth using a DataLoader with num_workers and prefetch_factor. The maximum speed will be your internet bandwidth (unless your CPU is a bottlenbeck for CPU operations like image decoding).

> > If we have already loaded the dataset, why doing to_iterable_dataset? Does it go through the dataset faster than map-style dataset?
> 
> yes, it makes a faster DataLoader for example (otherwise DataLoader uses `__getitem__` which is slower than iterating)

Okay, but `__getitem__` seems suitable for distributed settings. A distributed sampler would dispatch distinct indexes to each rank (rank0 got 0,1,2,3, rank1 got 4,5,6,7), however, if we make it `to_iterable_dataset`, then each rank needs to iterate all the samples, making it slower (i,e, rank1 got 0,1,2,3, rank2 got 0,1,2,3,(4,5,6,7))

What's your opinion here?

> however, if we make it to_iterable_dataset, then each rank needs to iterate all the samples, making it slower (i,e, rank1 got 0,1,2,3, rank2 got 0,1,2,3,(4,5,6,7))

Actually if you specify `to_iterable_dataset(num_shards=world_size)` (or a factor of world_size) and use a `torch.utils.data.DataLoader` then each rank will get a subset of the data thanks to the sharding. E.g. rank0 gets 0,1,2,3 and rank1 gets 4,5,6,7.

This is because `datasets.IterableDataset` subclasses `torch.utils.data.IterableDataset` and is aware of the current rank.

Got it, very nice features `num_shards` 👍🏻 

I would benchmark `to_iterable_dataset(num_shards=world_size)` against traditional map-style one in distributed settings in the near future.

Hi @lhoestq , I run a test for the speed in single node. Things are not expected as you mentioned before.

```python
import time

import datasets
from torch.utils.data import DataLoader


def time_decorator(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"Time taken: {end_time - start_time} seconds")
        return result

    return wrapper


dataset = datasets.load_dataset(
    "parquet", data_dir="my_dir", split="train"
)


@time_decorator
def load_dataset1():
    for _ in dataset:
        pass


@time_decorator
def load_dataloader1():
    for _ in DataLoader(dataset, batch_size=100, num_workers=5):
        pass


@time_decorator
def load_dataset2():
    for _ in dataset.to_iterable_dataset():
        pass


@time_decorator
def load_dataloader2():
    for _ in DataLoader(dataset.to_iterable_dataset(num_shards=5), batch_size=100, num_workers=5):
        pass


load_dataset1()
load_dataloader1()
load_dataset2()
load_dataloader2()
```
```bash
Resolving data files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 53192/53192 [00:00<00:00, 227103.16it/s]
Time taken: 100.36162948608398 seconds
Time taken: 70.09702134132385 seconds
Time taken: 343.09229612350464 seconds
Time taken: 132.8996012210846 seconds
```

1. Why `for _ in dataset.to_iterable_dataset()` is much slower than `for _ in dataset`
2. The `70 < 132`, the dataloader is slower when `to_iterable_dataset`

Loading in batches is faster than one example at a time. In your test the dataset is loaded in batches while the iterable_dataset is loaded one example at a time and the dataloader has a buffer to turn the examples to batches.

can you try this ?

```
batched_dataset = dataset.batch(100, num_proc=5)

@time_decorator
def load_dataloader3():
    for _ in DataLoader(batched_dataset.to_iterable_dataset(num_shards=5), batch_size=None, num_workers=5):
        pass
```

To be fair, I test the time including batching:
```python
@time_decorator
def load_dataloader3():
    for _ in DataLoader(dataset.batch(100, num_proc=5).to_iterable_dataset(num_shards=5), batch_size=None, num_workers=5):
        pass
```

```bash
Time taken: 49.722447633743286 seconds
```

I run another test about shuffling.

```python
@time_decorator
def load_map_dataloader1():
    for _ in DataLoader(dataset, batch_size=100, num_workers=5, shuffle=True):
        pass

@time_decorator
def load_map_dataloader2():
    for _ in DataLoader(dataset.batch(100, num_proc=5), batch_size=None, num_workers=5, shuffle=True):
        pass


@time_decorator
def load_iter_dataloader1():
    for _ in DataLoader(dataset.batch(100, num_proc=5).to_iterable_dataset(num_shards=5).shuffle(buffer_size=1000), batch_size=None, num_workers=5):
        pass

load_map_dataloader1()
load_map_dataloader2()
load_iter_dataloader1()
```

```bash
Time taken: 43.8506863117218 seconds
Time taken: 38.02591300010681 seconds
Time taken: 53.38815689086914 seconds
```


- What if I have custom collate_fn when batching?

- And if I want to shuffle the dataset, what's the correct order for `to_iterable_dataset(num_shards=x)`, `batch()` and `shuffle()`. Is `dataset.batch().to_iterable_dataset().shuffle()`? This is not faster than map-style dataset