Batched IterableDataset
State: open

### Feature request

Hi, 

could you add an implementation of a batched `IterableDataset`. It already support an option to do batch iteration via `.iter(batch_size=...)` but this cannot be used in combination with a torch `DataLoader` since it just returns an iterator.

### Motivation

The current implementation loads each element of a batch individually which can be very slow in cases of a big batch_size. I did some experiments [here](https://discuss.huggingface.co/t/slow-dataloader-with-big-batch-size/57224) and using a batched iteration would speed up data loading significantly.

### Your contribution

N/A

This is exactly what I was looking for. It would also be very useful for me :-)

This issue is really smashing the selling point of HF datasets... The only workaround I've found so far is to create a customized IterableDataloader which improves the loading speed to some extent.

For example I've a HF dataset `dt_train` with `len(dt_train) == 1M`. Using plain DataLoader is extremely slow:
```
%%time
dl_train = DataLoader(dt_train, batch_size=128, shuffle = True)
for batch in dl_train:
    pass
``` 

```
CPU times: user 24min 35s, sys: 704 ms, total: 24min 36s
Wall time: 24min 37s
```
And DataLoader works even worse with HF's iterable_dataset:
```
%%time
dt_train_ = dt_train.with_format(None).to_iterable_dataset(num_shards=64).shuffle(buffer_size=10_000)
dl_train = DataLoader(dt_train_, batch_size=128)
for batch in dl_train:
    pass
```
```
CPU times: user 1h 6min 2s, sys: 4.28 s, total: 1h 6min 6s
Wall time: 1h 7min 53s
```
Workaround by running a customized wrapper:
```
%%time
from torch.utils.data import DataLoader, IterableDataset

class Dataset2Iterable(IterableDataset):
    """
    Wrapper to use a HF dataset as pytorch IterableDataset to speed up data loading.
    """
    def __init__(self, dataset, batch_size=1, shuffle=True):
        super(Dataset2Iterable).__init__()
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self):
        if self.shuffle: self.dataset.shuffle()
        return self.dataset.iter(batch_size=self.batch_size)

dl_train = DataLoader(Dataset2Iterable(dt_train, batch_size = 128), batch_size=1, num_workers=0)
for n in range(2):
    for batch in dl_train:
        pass
```
The speed still is slower than using tensorflow's loader but improved a lot than previous code:
```
CPU times: user 4min 18s, sys: 0 ns, total: 4min 18s
Wall time: 4min 20s
```
Note that the way I implemented `Dataset2Iterable` will only work with `num_workers=0`.

I can confirm that @zhh210's solution works with `num_workers=0`. However, for my use case, this was still slower than tokenizing on the fly through a collator and leveraging multiple workers in the dataloder.

@lhoestq I think this is an important use case (e.g., streaming from a large dataset, online or stored on disk). What do you think might be the best solution to move forward?

I guess it can be implemented using a batched`.map()` under the hood that returns a single item containing the input batch.

In the meantime you can use this:

```python
def batch(unbatched: dict[str, list]) -> dict[str, list]:
    return {k: [v] for k, v in unbatched}

batched_dataset = dataset.map(batch, batched=True, batch_size=batch_size)
```

Though it would be great to have a `.batch()` method indeed, I'd be happy to help with anyone wants to open a PR

If no one else is planning to work on this, I can take it on. I'll wait until next week, and if no one has started a PR by then, I'll go ahead and open one.

It looks like the implementation of IterableDataset is still using a hardcoded batch size of 1. For example in line 2063 in [`/datsets/src/datasets/iterable_dataset.py`](https://github.com/huggingface/datasets/blob/3.0.2/src/datasets/iterable_dataset.py#L2063). Iterating over IterableDataset with large batch sizes therefore remains slow, even when using `batch()`. I guess then the data are not being read from one contiguous chunk of memory. Instead every example is retrieved one by one, leading to long dataloading times. As a minimal example: Load c4 dataset and iterate over it with a large batch size.
```python
import datasets
from timeit import default_timer as timer
c4 = datasets.load_dataset("allenai/c4", "en", streaming=True, split="train")
c4_batched = c4.batch(512**2) # use large batch size
iterator = iter(c4_batched)
for i in range(5):
    start_time=timer()
    next(iterator) # get next batch
    end_time = timer()
    print(f"time for one batch: {end_time-start_time}")
```
This results in the following output for me:
time for one batch: 12.615376660600305
time for one batch: 13.011422813870013
time for one batch: 14.157325950451195
time for one batch: 14.225894245319068
time for one batch: 13.898222777992487

Because I want to use my IterableDataset with the pytorch dataloader I rewrote the `__iter_pytorch__` and the `__iter__` functions like so and am getting much faster dataloading times. I marked the lines I changed with "# changed here":

```python
from datasets.iterable_dataset import _convert_to_arrow
from datasets.formatting import TensorFormatter, get_formatter
from datasets.features.features import cast_to_python_objects
import sys
import fsspec.asyn
from itertools import islice
from datasets.utils.logging import get_logger
from datasets.iterable_dataset import _examples_to_batch, _apply_feature_types_on_batch, _apply_feature_types_on_example

logger = get_logger(__name__)

def __iter__(self):
    if "torch" in sys.modules:
        import torch.utils.data

        worker_info = torch.utils.data.get_worker_info()
        if isinstance(self, torch.utils.data.IterableDataset) and worker_info is not None:
            # We're a torch.utils.data.IterableDataset in a PyTorch worker process
            yield from self._iter_pytorch()
            return

    ex_iterable = self._prepare_ex_iterable_for_iteration(batch_size=self.batch_size, drop_last_batch=self.drop_last_batch) # changed here
    if self._formatting:
        formatter = get_formatter(self._formatting.format_type, features=self.features)
        format_dict = (
            formatter.recursive_tensorize if isinstance(formatter, TensorFormatter) else cast_to_python_objects
        )
    else:
        format_dict = None

    if self._formatting and (ex_iterable.iter_arrow or self._formatting.format_type == "arrow"):
        if ex_iterable.iter_arrow:
            iterator = ex_iterable.iter_arrow()
        else:
            iterator = _convert_to_arrow(ex_iterable, batch_size=self.batch_size) # changed here
        for key, pa_table in iterator:
            yield formatter.format_row(pa_table)
        return

    for key, example in ex_iterable:
        if self.features and not ex_iterable.is_typed:
            # `IterableDataset` automatically fills missing columns with None.
            # This is done with `_apply_feature_types_on_example`.
            example = _apply_feature_types_on_example(
                example, self.features, token_per_repo_id=self._token_per_repo_id
            )
        yield format_dict(example) if format_dict else example



def _iter_pytorch(self):
    ex_iterable = self._prepare_ex_iterable_for_iteration(batch_size=self.batch_size, drop_last_batch=self.drop_last_batch) # changed here
    # Fix for fsspec when using multiprocess to avoid hanging in the ML training loop. (only required for fsspec >= 0.9.0)
    # See https://github.com/fsspec/gcsfs/issues/379
    fsspec.asyn.reset_lock()
    # check if there aren't too many workers
    import torch.utils.data

    worker_info = torch.utils.data.get_worker_info()
    if self._is_main_process() and ex_iterable.n_shards < worker_info.num_workers:
        logger.warning(
            f"Too many dataloader workers: {worker_info.num_workers} (max is dataset.n_shards={ex_iterable.n_shards}). "
            f"Stopping {worker_info.num_workers - ex_iterable.n_shards} dataloader workers."
        )
        logger.info(
            f"To parallelize data loading, we give each process some shards (or data sources) to process. "
            f"Therefore it's unnecessary to have a number of workers greater than dataset.n_shards={ex_iterable.n_shards}. "
            f"To enable more parallelism, please split the dataset in more files than {ex_iterable.n_shards}."
        )
    # split workload
    _log_prefix = f"node#{self._distributed.rank} " if self._distributed else ""
    shards_indices = ex_iterable.split_shard_indices_by_worker(worker_info.id, worker_info.num_workers)
    if shards_indices:
        logger.debug(
            f"{_log_prefix}dataloader worker#{worker_info.id}, ': Starting to iterate over {len(shards_indices)}/{ex_iterable.n_shards} shards."
        )
        ex_iterable = ex_iterable.shard_data_sources(worker_id=worker_info.id, num_workers=worker_info.num_workers)
        self._state_dict = ex_iterable._init_state_dict()
        if self._starting_state_dict:
            ex_iterable.load_state_dict(self._starting_state_dict)

        if self._formatting:
            formatter = get_formatter(self._formatting.format_type, features=self.features)
            format_dict = (
                formatter.recursive_tensorize if isinstance(formatter, TensorFormatter) else cast_to_python_objects
            )
        else:
            format_dict = None

        if self._formatting and (ex_iterable.iter_arrow or self._formatting == "arrow"):
            if ex_iterable.iter_arrow:
                iterator = ex_iterable.iter_arrow()
            else:
                iterator = _convert_to_arrow(ex_iterable, batch_size=self.batch_size) # changed here
            if self.batch_size > 1: # changed here until end of file
                for key, pa_table in iterator:
                    yield formatter.format_batch(pa_table)
                return
            else:
                for key, pa_table in iterator:
                    yield formatter.format_row(pa_table)
                return

        iterator = iter(ex_iterable)
        if self.batch_size > 1:
            for key, example in iterator:
                    # If batched, first build the batch
                    examples = [example] + [example for key, example in islice(iterator, self.batch_size - 1)]
                    if self.drop_last_batch and len(examples) < self.batch_size:  # ignore last batch
                        return
                    batch = _examples_to_batch(examples)
                    if self.features and not ex_iterable.is_typed:
                        # `IterableDataset` automatically fills missing columns with None.
                        # This is done with `_apply_feature_types_on_batch`.
                        batch = _apply_feature_types_on_batch(batch, self.features, token_per_repo_id=self._token_per_repo_id)
                    yield format_dict(batch) if format_dict else batch
        else:
            for key, example in ex_iterable:
                if self.features and not ex_iterable.is_typed:
                    # `IterableDataset` automatically fills missing columns with None.
                    # This is done with `_apply_feature_types_on_example`.
                    example = _apply_feature_types_on_example(
                        example, self.features, token_per_repo_id=self._token_per_repo_id
                    )
                yield format_dict(example) if format_dict else example
        logger.debug(
            f"{_log_prefix}dataloader worker#{worker_info.id}, ': Finished iterating over {len(shards_indices)}/{ex_iterable.n_shards} shards."
        )
    else:
        logger.debug(
            f"{_log_prefix}dataloader worker#{worker_info.id}, ': Stopping... Number of dataset shards < num_workers ({ex_iterable.n_shards}<{worker_info.num_workers})."
        )
```

For anyone wanting to try it you can patch it into datasets by overwriting the function via `setattr(datasets.IterableDataset, '_iter_pytorch', _iter_pytorch)`

I don't really know what most of the rest of the code is doing so no idea if this is a valid fix or not, but it seems to work for me.
Example of running the fix:
```python
from torch.utils.data.dataloader import DataLoader
c4.batch_size = 512**2 # set batch size here
dataloader = Dataloader(c4, batch_size=None) # use custom batching from IterableDataset
iterator = iter(dataloader)
for i in range(5):
    start_time=timer()
    next(iterator) #get the batch
    end_time = timer()
    print(f"time for one batch: {end_time-start_time}")
```
I now get 
time for one batch: 0.6047679269686341
time for one batch: 0.000248616561293602
time for one batch: 0.00017435848712921143
time for one batch: 0.00015910807996988297
time for one batch: 0.00015317369252443314

I love the datasets library and it would be great if iterating with large batch sizes would be supported directly, either with a similar fix to mine or in some other way :)

Hi @taczin , thanks for reporting !

Indeed the `IterableDataset.batch()` implementation is quite naive is manipulates python objects:

https://github.com/huggingface/datasets/blob/d37ed46ebf45981131bd3678173dbb4b7e2b2f1a/src/datasets/iterable_dataset.py#L3026-L3029

However it can be much faster if it can be applied on the Arrow data, maybe using something like this (untested)

```python
def batch_fn(unbatched): 
    return {k: [v] for k, v in unbatched.items()} 

def batch_fn_arrow(unbatched_pa_table): 
    offsets = pa.array([0, len(unbatched_pa_table)])
    return pa.Table.from_arrays([
        pa.ListArray.from_arrays(offsets, unbatched_pa_table[k])
        for k in unbatched_pa_table.column_names
    ], unbatched_pa_table.column_names)

if self._ex_iterabe.iter_arrow:
    return self.with_format("arrow").map(
        batch_fn_arrow, batched=True, batch_size=batch_size, drop_last_batch=drop_last_batch
    ).with_format(self._formatting.format_type if self._formatting else None)
else:
    return self.map(batch_fn, batched=True, batch_size=batch_size, drop_last_batch=drop_last_batch)
```

Hi @lhoestq , thanks for your answer. I was wondering: is there a reason why the internal call to `ex_iterable = self._prepare_ex_iterable_for_iteration()` in the IterableDataset code does not pass the batch size even though it could? If not passed the default of 1 is used, leading to the observed slow loading.

After calling `.batch()`, `_prepare_ex_iterable_for_iteration` should use batch_size=1 since now each row in the dataset is actually a batch of the original dataset.