Concatenate IterableDataset instances and distribute underlying shards in a RoundRobin manner

### Feature request

I would like to be able to concatenate multiple `IterableDataset` with possibly different features. I would like to then be able to stream the results in parallel (both using DDP and multiple workers in the pytorch DataLoader). I want the merge of datasets to be well balanced between the different processes.

### Motivation

I want to train a model on a combination of datasets, which I can convert to a single representation. This applies to converting different datasets items to the same Python class, as using a tokenizer on multiple modalities.

Assuming that my original datasets are not necessarily well balanced as they may have different size and thus different number of shards, I would like the merged dataset to be distributed evenly over the multiple processes. I don't mind if it's not perfectly balanced, and as result, some workers of the torch DataLoader do nothing, as long as the DDP is properly handled causing no deadlock.

### What I've tried

I've tried the two functions already provided in datasets, namely `interleave_datasets` and `concatenate_datasets`. 

- Interleave seems to be the best approach of what I'm trying to do. However, it doesn't suit my purpose because as I understand it, it stops as soon as one of the dataset source is exhausted, or repeat the smallest source items until the largest is exhausted. I would like something in-between, similarly to what [roundrobin does](https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.roundrobin).
- Concatenate does not mix the data enough and one dataset may be overrepresented in some early batches.

Let's consider we have 3 datasets composed of different number of shards as follow [[s0_0, s0_1], [s1_0], [s2_0, s2_1, s2_3]], where s denotes the underlying shard, the first index the dataset and the second the shard number.
If we request 3 shards in the `shard_data_source` we should obtain the following:
index 0 gets s0_0 s2_0
index 1 gets s0_1 s2_1
index 2 gets s1_0 s2_3

I started implementing the following, but I'm afraid my sharding logic is incorrect.

```python
from copy import deepcopy
from itertools import chain, islice

import datasets
import numpy as np
from datasets import IterableDataset
from datasets.iterable_dataset import _BaseExamplesIterable
from more_itertools import roundrobin


class MixMultiSourcesExampleIterable(_BaseExamplesIterable):
    def __init__(self, ex_iterables: list[_BaseExamplesIterable]):
        super().__init__()
        self.ex_iterables = ex_iterables

    def _init_state_dict(self) -> dict:
        self._state_dict = {
            "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
            "type": self.__class__.__name__,
        }
        return self._state_dict

    @property
    def num_shards(self) -> int:
        return sum(ex_iterable.num_shards for ex_iterable in self.ex_iterables)

    def __iter__(self):
        yield from roundrobin(*self.ex_iterables)

    def shuffle_data_sources(self, generator: np.random.Generator) -> "MixMultiSourcesExampleIterable":
        """Shuffle the list of examples iterable, as well as each underlying examples iterable."""
        rng = deepcopy(generator)
        ex_iterables = list(self.ex_iterables)
        rng.shuffle(ex_iterables)
        ex_iterables = [ex_iterable.shuffle_data_sources(generator) for ex_iterable in ex_iterables]
        return MixMultiSourcesExampleIterable(ex_iterables)

    def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "MixMultiSourceExampleIterable":
        """Shard the underlying iterables in a roundrobin manner.

        Let's consider we have our iterables as [[s0_0, s0_1], [s1_0], [s2_0, s2_1, s2_3]],
        and we request 3 shards.
        index 0 gets s0_0 s2_0
        index 1 gets s0_1 s2_1
        index 2 gets s1_0 s2_3
        """
        return MixMultiSourcesExampleIterable(
            list(
                islice(
                    # flatten all underlying iterables
                    chain.from_iterable([ex_iterable.shard_data_sources(1, 0) for ex_iterable in self.ex_iterables]),
                    # offset the starting point by the index
                    index,
                    # take over the full list, so exhaust the iterators
                    None,
                    # step by the number of shards requested
                    num_shards,
                )
            )
        )


def mix_dataset(iterable_datasets: list[datasets.IterableDataset]) -> IterableDataset:
    ex_iterable = MixMultiSourcesExampleIterable([ds._ex_iterable for ds in iterable_datasets])
    return IterableDataset(
        ex_iterable, distributed=iterable_datasets[0]._distributed, formatting=iterable_datasets[0]._formatting
    )

```

### Questions

- Am I missing something? Is there a way to use `interleave_datasets` or `concatenate_datasets` to fit my purpose?
- Would it be the right approach to spread the maximum number of underlying shards across my different processes?

### Your contribution

As much as I can.

Comments:
# With `datasets.Dataset`

Here is an small script that shows the distribution differences of samples between `interleave_datasets`, `concatenate_datasets` and `concatenate_datasets` + shuffling.

```python
import datasets as hf_datasets

def gen(dataset: int, n_samples: int):
    for i in range(n_samples):
        yield {"dataset": dataset, "sample": i}

ds_1 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 0, "n_samples": 2})
ds_2 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 1, "n_samples": 1})
ds_3 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 2, "n_samples": 3})

n_workers = 3
print(f"Simulate run with {n_workers} workers")

print("Interleave datasets")
for w in range(n_workers):
    ds_interleave = hf_datasets.interleave_datasets([ds_1, ds_2, ds_3]).shard(n_workers, w)
    for i, sample in enumerate(ds_interleave):
        print(f"Worker {w} process sample {i} {sample}")

print("Concatenate datasets")
for w in range(n_workers):
    ds_concatenate = hf_datasets.concatenate_datasets([ds_1, ds_2, ds_3]).shard(n_workers, w)
    for i, sample in enumerate(ds_concatenate):
        print(f"Worker {w} process sample {i} {sample}")

print("Concated and shuffled datasets")
for w in range(n_workers):
    ds_concatenate = hf_datasets.concatenate_datasets([ds_1, ds_2, ds_3]).shuffle().shard(n_workers, w)
    for i, sample in enumerate(ds_concatenate):
        print(f"Worker {w} process sample {i} {sample}")
```

> Interleave datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 1 process sample 0 {'dataset': 1, 'sample': 0}
Worker 2 process sample 0 {'dataset': 2, 'sample': 0}

> Concatenate datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 0 process sample 1 {'dataset': 0, 'sample': 1}
Worker 1 process sample 0 {'dataset': 1, 'sample': 0}
Worker 1 process sample 1 {'dataset': 2, 'sample': 0}
Worker 2 process sample 0 {'dataset': 2, 'sample': 1}
Worker 2 process sample 1 {'dataset': 2, 'sample': 2}

> Concated and shuffled datasets
Worker 0 process sample 0 {'dataset': 2, 'sample': 2}
Worker 0 process sample 1 {'dataset': 2, 'sample': 0}
Worker 1 process sample 0 {'dataset': 0, 'sample': 1}
Worker 1 process sample 1 {'dataset': 2, 'sample': 1}
Worker 2 process sample 0 {'dataset': 2, 'sample': 2}
Worker 2 process sample 1 {'dataset': 0, 'sample': 0}

Without shuffling, round robin would yield:
> Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 0 process sample 1 {'dataset': 2, 'sample': 0}
Worker 1 process sample 0 {'dataset': 0, 'sample': 1}
Worker 1 process sample 1 {'dataset': 2, 'sample': 1}
Worker 2 process sample 0 {'dataset': 1, 'sample': 0}
Worker 2 process sample 1 {'dataset': 2, 'sample': 2}

# With `datasets.IterableDataset`

The above works for `Dataset`, but with a sharded `IterableDataset` some data get discarded. See the following results obtained with the script below.

> Simulate run with 3 workers

> Interleave datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 1 fails with list index out of range.
Worker 2 fails with list index out of range.
With dataloader
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
{'dataset': tensor([0]), 'sample': tensor([0])}

> Concatenate datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 0 process sample 1 {'dataset': 1, 'sample': 0}
Worker 0 process sample 2 {'dataset': 2, 'sample': 0}
Worker 1 fails with list index out of range
Worker 2 fails with list index out of range
With dataloader
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
{'dataset': tensor([0]), 'sample': tensor([0])}
{'dataset': tensor([1]), 'sample': tensor([0])}
{'dataset': tensor([2]), 'sample': tensor([0])}

> Concated and shuffled datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 0 process sample 1 {'dataset': 1, 'sample': 0}
Worker 0 process sample 2 {'dataset': 2, 'sample': 0}
Worker 1 fails with list index out of range
Worker 2 fails with list index out of range
With dataloader
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
{'dataset': tensor([0]), 'sample': tensor([0])}
{'dataset': tensor([1]), 'sample': tensor([0])}
{'dataset': tensor([2]), 'sample': tensor([0])}

<details>

<summary>Experiment script</summary>

```python
ds_1 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 0, "n_samples": 2}).to_iterable_dataset(
    num_shards=2
)
ds_2 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 1, "n_samples": 1}).to_iterable_dataset(
    num_shards=1
)
ds_3 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 2, "n_samples": 3}).to_iterable_dataset(
    num_shards=3
)

n_workers = 3
print(f"Simulate run with {n_workers} workers")

print("\nInterleave datasets")
ds_interleave = hf_datasets.interleave_datasets([ds_1, ds_2, ds_3])
for w in range(n_workers):
    try:
        for i, sample in enumerate(ds_interleave.shard(n_workers, w)):
            print(f"Worker {w} process sample {i} {sample}")
    except IndexError as e:
        print(f"Worker {w} fails with {e}.")

print("With dataloader")
for sample in torch.utils.data.DataLoader(ds_interleave, num_workers=n_workers):
    print(f"{sample}")

print("\nConcatenate datasets")
ds_concatenate = hf_datasets.concatenate_datasets([ds_1, ds_2, ds_3])
for w in range(n_workers):
    try:
        for i, sample in enumerate(ds_concatenate.shard(n_workers, w)):
            print(f"Worker {w} process sample {i} {sample}")
    except IndexError as e:
        print(f"Worker {w} fails with {e}")

print("With dataloader")
for sample in torch.utils.data.DataLoader(ds_concatenate, num_workers=n_workers):
    print(f"{sample}")

print("\nConcated and shuffled datasets")
ds_concatenate = hf_datasets.concatenate_datasets([ds_1, ds_2, ds_3]).shuffle()
for w in range(n_workers):
    try:
        for i, sample in enumerate(ds_concatenate.shard(n_workers, w)):
            print(f"Worker {w} process sample {i} {sample}")
    except IndexError as e:
        print(f"Worker {w} fails with {e}")

print("With dataloader")
for sample in torch.utils.data.DataLoader(ds_concatenate, num_workers=n_workers):
    print(f"{sample}")
```

</details>

# Round Robin with fixed logic

> I started implementing the following, but I'm afraid my sharding logic is incorrect.

Here is a solution for mixing the data in a round robin fashion that allows to distribute the data to all workers. In the previous example above only 1 worker over 3 was actually retrieving data, which resulted in discarding some data.

```python
def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "MixMultiSourceExampleIterable":
    """Shard the underlying iterables in a roundrobin manner.

    Let's consider we have our iterables as [[s0_0, s0_1], [s1_0], [s2_0, s2_1, s2_3]],
    and we request 3 shards.
    index 0 gets s0_0 s2_0
    index 1 gets s0_1 s2_1
    index 2 gets s1_0 s2_3
    """
    return MixMultiSourcesExampleIterable(
        list(
            islice(
                # flatten all underlying iterables (fixed logic)
                [
                    ex_iterable.shard_data_sources(ex_iterable.num_shards, index)
                    for ex_iterable in self.ex_iterables
                    for index in range(ex_iterable.num_shards)
                ],
                # offset the starting point by the index
                index,
                # take over the full list, so exhaust the iterators
                None,
                # step by the number of shards requested
                num_shards,
            )
        )
    )
```

Editing the example above with the following we obtain the expected result:
```python
print("\nMix datasets")
ds_mix = mix_dataset([ds_1, ds_2, ds_3])
for w in range(n_workers):
    try:
        for i, sample in enumerate(ds_mix.shard(n_workers, w)):
            print(f"Worker {w} process sample {i} {sample}")
    except IndexError as e:
        print(f"Worker {w} fails with {e}")

print("With dataloader")
for sample in torch.utils.data.DataLoader(ds_mix, num_workers=n_workers):
    print(f"{sample}")
```
> Mix datasets
Mix datasets
Worker 0 process sample 0 {'dataset': 0, 'sample': 0}
Worker 0 process sample 1 {'dataset': 2, 'sample': 0}
Worker 1 process sample 0 {'dataset': 0, 'sample': 1}
Worker 1 process sample 1 {'dataset': 2, 'sample': 1}
Worker 2 process sample 0 {'dataset': 1, 'sample': 0}
Worker 2 process sample 1 {'dataset': 2, 'sample': 2}
With dataloader
{'dataset': tensor([0]), 'sample': tensor([0])}
{'dataset': tensor([0]), 'sample': tensor([1])}
{'dataset': tensor([1]), 'sample': tensor([0])}
{'dataset': tensor([2]), 'sample': tensor([0])}
{'dataset': tensor([2]), 'sample': tensor([1])}
{'dataset': tensor([2]), 'sample': tensor([2])}

# Questions 

- The example is quite small, showing that some data get discarded, but on large datasets is this significant?
- How does the suggested solution interplays with shuffling?






# Larger Experiment

> The example is quite small, showing that some data get discarded, but on large datasets is this significant?

Continuing the experiment above, but with 3 larger and unbalanced datasets, with respectively 1000, 150, and 300 samples, and a dataloader with 4 workers:
 
> Interleave datasets
With dataloader
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
Yield 300 samples

> Concatenate datasets
With dataloader
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
Yield 705 samples

> Concated and shuffled datasets
With dataloader
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
Yield 705 samples

> Mix datasets
With dataloader
Yield 1405 samples

The dataset mixing proposed above is the only one that yields all the samples while using all the dataloaders.
Additional checks should include training metrics (does it improve training quality to mix the data like this), and behavior check in a DDP settings, we don't want to face any deadlock due to some GPU having more batches than other. But this later point should be already handled by the iterator of the `IterableDataset`.

# Follow up?

@lhoestq would there be any interest in making a PR of it? Otherwise I can close the issue as I found a solution to my problem. 

I believe this PR could solve your issue? :)

https://github.com/huggingface/datasets/pull/7786

> I believe this PR could solve your issue? :)

Thank you @lhoestq for the reply.
I have just tested it with the script above. It gives:

> Interleave datasets without replacement
With dataloader
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
Yield 705 samples

If we compare with the original `interleave_dataset` method it produces 405 samples more. However, it only uses 1 worker on the 4 available. Moreover it doesn't yield all the samples as the mixing strategy with RoundRobin above does (1405 samples vs 705).

@LTMeyer With the following script and using the code from #7786 I get all 1450 samples

```
import datasets as hf_datasets


def gen(dataset: int, n_samples: int):
    for i in range(n_samples):
        yield {"dataset": dataset, "sample": i}


ds_1 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 0, "n_samples": 1000}).to_iterable_dataset()
ds_2 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 1, "n_samples": 150}).to_iterable_dataset()
ds_3 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 2, "n_samples": 300}).to_iterable_dataset()

print("Interleave datasets")
ds_interleave = hf_datasets.interleave_datasets(
    [ds_1, ds_2, ds_3],
    probabilities=[1 / 3, 1 / 3, 1 / 3],
    stopping_strategy="all_exhausted_without_replacement",
)
for i, sample in enumerate(ds_interleave):
    print(f"process sample {i} {sample}")
```
I'm not sure on the workers side how many will be spawned and so on. 

> [@LTMeyer](https://github.com/LTMeyer) With the following script and using the code from [#7786](https://github.com/huggingface/datasets/pull/7786) I get all 1450 samples

This depends on the number of shards and the number of processes being used.
In the example below there is only one shard per dataset (the default of `to_iterable_dataset` method). Then, the for loop is running in the main process. It thus consumes all the shards, hence the 1450 samples.

> 
> ```
> import datasets as hf_datasets
> 
> 
> def gen(dataset: int, n_samples: int):
>     for i in range(n_samples):
>         yield {"dataset": dataset, "sample": i}
> 
> 
> ds_1 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 0, "n_samples": 1000}).to_iterable_dataset()
> ds_2 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 1, "n_samples": 150}).to_iterable_dataset()
> ds_3 = hf_datasets.Dataset.from_generator(gen, gen_kwargs={"dataset": 2, "n_samples": 300}).to_iterable_dataset()
> 
> print("Interleave datasets")
> ds_interleave = hf_datasets.interleave_datasets(
>     [ds_1, ds_2, ds_3],
>     probabilities=[1 / 3, 1 / 3, 1 / 3],
>     stopping_strategy="all_exhausted_without_replacement",
> )
> for i, sample in enumerate(ds_interleave):
>     print(f"process sample {i} {sample}")
> ```
> 


> I'm not sure on the workers side how many will be spawned and so on.

While using the data to train a model, I would like to use the `torch.utils.data.DataLoader` to feed batches of data to my model. To make the data loading fast, it is common to use `num_workers>0` in the dataloader. This will consume data in parallel. In practice, it copies the dataset instance and read in parallel different chunks of data. These chunks correspond to the underlying shards of the iterable dataset.

If we have 1 shard per dataset, as it is the case in the example above, the dataloading will indeed get all the 1450 samples, but it will run only in one process even if multiple are available. This is inefficient because it doesn't utilize all available resources. See the script and results below.

```python
for num_workers in [0, 1, 2, 3, 4]:
    print(f"Dataloader with {num_workers} workers.")
    dataloader = DataLoader(ds_interleave, num_workers=num_workers, batch_size=1)
    for i, sample in enumerate(dataloader, start=1):
       pass
    print(f"{i} processed samples")
```

```
Dataloader with 0 workers.
1450 processed samples
Dataloader with 1 workers.
1450 processed samples
Dataloader with 2 workers.
Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
1450 processed samples
Dataloader with 3 workers.
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
1450 processed samples
Dataloader with 4 workers.
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
1450 processed samples
```

Now if we shard our data differently, like 2, 1, and 3 for each dataset respectively as the [previous example](https://github.com/huggingface/datasets/issues/7792#issuecomment-3345970293), and use a dataloader with different number of workers (same script as above), we obtain:

```
Dataloader with 0 workers.
1450 processed samples
Dataloader with 1 workers.
1450 processed samples
Dataloader with 2 workers.
Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
850 processed samples
Dataloader with 3 workers.
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
750 processed samples
Dataloader with 4 workers.
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
750 processed samples
```

I added a small fix to your PR @radulescupetru to try to make @LTMeyer 's example work :)

Can you confirm it works for you now @LTMeyer ?

Note that maximum parallelism requires each subset to have num_shards >= num_workers, otherwise there aren't enough shards to distribute to every worker for interleaving. In your example one of the subsets has only 1 shard, so only 1 worker can take care of interleaving.

> Can you confirm it works for you now [@LTMeyer](https://github.com/LTMeyer) ?

Result with https://github.com/huggingface/datasets/pull/7786/commits/a547d81469128bea4acc3bcc2a4a6a95968936ee:
```
Dataloader with 0 workers.
1450 processed samples
Dataloader with 1 workers.
1450 processed samples
Dataloader with 2 workers.
Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
1450 processed samples
Dataloader with 3 workers.
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
1450 processed samples
Dataloader with 4 workers.
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
1450 processed samples
```

 I have checked with the script above and I confirm that all samples are now correctly returned, thank you @lhoestq .

> Note that maximum parallelism requires each subset to have num_shards >= num_workers, otherwise there aren't enough shards to distribute to every worker for interleaving. In your example one of the subsets has only 1 shard, so only 1 worker can take care of interleaving.

This point I'm not sure I understand. That is maybe where @radulescupetru's intent and mine differ. Why should we limit the number of workers to the minimum number of shards? My initial goal was to distribute shards among workers to maximize data loading speed, and to mix the data so batches are representative of the whole dataset and diverse enough (hence the round-robin). 

In the example above, we have 6 shards in total, can we not distribute these shards among workers? That what the `MixMultiSourcesExampleIterable` in https://github.com/huggingface/datasets/issues/7792#issuecomment-3345970293 above does.
- If 2 workers, 3 shards for each. 
- If 3 workers, 2 shards for each.
- If 4 workers, the 2 first ones get 2 shards while the two last ones get only 1.
- Above 6 workers, the 6 first ones get 1 shard each, and the remaining workers get none.




@LTMeyer I think it's just a design choice that datasets library took. From my interaction with it, it seems that even when concatenating or interleaving, individual components are still treated individually (for example, num_shards is not summed).

I guess in a real scenario you wouldn't  end up with 1 shard only, but it's true that you need to be a bit careful with the setup. For workers it's a bit more automated in the sense that if you have more it will stop the extra ones, but when distributing a dataset over multiple gpus it's even more tricky as if the number of shards is not a factor of world size iterating is slower.

> [@LTMeyer](https://github.com/LTMeyer) I think it's just a design choice that datasets library took. From my interaction with it, it seems that even when concatenating or interleaving, individual components are still treated individually (for example, num_shards is not summed).

Indeed. I am curious to know if there is any explanation for this choice that I am missing.

> I guess in a real scenario you wouldn't end up with 1 shard only, but it's true that you need to be a bit careful with the setup. 

In my case I would like to mix many small datasets which are individually based on only few shards. So it's actually close to the case with 1 shard only.

> For workers it's a bit more automated in the sense that if you have more it will stop the extra ones, but when distributing a dataset over multiple gpus it's even more tricky as if the number of shards is not a factor of world size iterating is slower.

My understanding is that, in a multi-gpu settings, we want each GPU to receive the same number of batches to avoid deadlock in any synchronization process. 
Multi-GPU related sharding of the `IterableDataset` is managed there https://github.com/huggingface/datasets/blob/4.1.1/src/datasets/iterable_dataset.py#L2371-L2392,
while the sharding for dataloaders with multiple workers is handled there https://github.com/huggingface/datasets/blob/4.1.1/src/datasets/iterable_dataset.py#L2292-L2314.

Here is a script to check the behavior in case of multi-gpus, using `split_dataset_by_node`. In the example I consider just 2 GPUs.

```python
world_size = 2
for num_workers in [0, 1, 2, 3, 4]:
    for rank in range(world_size):
        print(f"Rank {rank}")
        ds_interleave_rank = split_dataset_by_node(ds_interleave, rank, world_size)
        print(f"Dataloader with {num_workers} workers.")
        dataloader = DataLoader(ds_interleave_rank, num_workers=num_workers, batch_size=1)
        for i in enumerate(dataloader, start=1):
            pass
        print(f"{i} processed samples")
    print("\n")
```

The results using https://github.com/huggingface/datasets/pull/7786/commits/455bfaaa6d574aa9d9c9592baee390017512cc5f:
```
Rank 0
Dataloader with 0 workers.
725 processed samples
Rank 1
Dataloader with 0 workers.
725 processed samples


Rank 0
Dataloader with 1 workers.
725 processed samples
Rank 1
Dataloader with 1 workers.
725 processed samples


Rank 0
Dataloader with 2 workers.
Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
725 processed samples
Rank 1
Dataloader with 2 workers.
725 processed samples


Rank 0
Dataloader with 3 workers.
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
725 processed samples
Rank 1
Dataloader with 3 workers.
725 processed samples


Rank 0
Dataloader with 4 workers.
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
725 processed samples
Rank 1
Dataloader with 4 workers.
725 processed samples
```

If now I use the mixing described above the results are:
```
Rank 0
Dataloader with 0 workers.
750 processed samples
Rank 1
Dataloader with 0 workers.
700 processed samples


Rank 0
Dataloader with 1 workers.
750 processed samples
Rank 1
Dataloader with 1 workers.
700 processed samples


Rank 0
Dataloader with 2 workers.
750 processed samples
Rank 1
Dataloader with 2 workers.
700 processed samples


Rank 0
Dataloader with 3 workers.
750 processed samples
Rank 1
Dataloader with 3 workers.
700 processed samples


Rank 0
Dataloader with 4 workers.
750 processed samples
Rank 1
Dataloader with 4 workers.
700 processed samples
```

Different GPUs received different number of batches which is problematic. The interleave method, on the other hand, feeds each GPU with the same number of batches. Nonetheless, it doesn't leverage all available workers.
I'll check if I can fix the distribution of shards across GPU in the last configuration.

When concatenating or interleaving, the resulting `num_shards` is the *minimum `num_shards` of the input datasets*. This allows each new shard to always contain data from every input dataset. This ensures in every shard the right sampling when interleaving and the right data order when concatenating.

Summing the dataset shards isn't ideal since each shard would contain data from only one of the dataset and would not contain any interleaved/concatenated data.

Thank you @lhoestq, it makes perfect sense. The part I am missing is that if I concatenate many datasets with small number of shards it will result in a global dataset with not so many shards, thus limiting the use of available workers. Data loading will be consequently inefficient. I was looking for a solution to leverage all parallelism available to maximize data loading speed.

My original use case was:
I want to use a dataset stored on the HF hub. It is composed of many subfolders. Each of this subfolder contain only a few shards. I would like to use the dataset but only on a subset of folders, while keeping information about the origin of each sample (i.e. from which subfolder they come from).
The first part would possible with the `data_files` argument of `load_dataset` method. However, I would not have the origin information about the sample, as it is not provided in the original dataset. I was thus thinking about considering each subfolder as an independent HF iterable dataset and concatenate them. This method does not work because it drastically reduces the dataloading efficiency due to the low number of shards.

> Summing the dataset shards isn't ideal `since` each shard would contain data from only one of the dataset and would not contain any interleaved/concatenated data.

This is not necessarily a problem for my use case. It will be the case for the original dataset anyway.

Also, I notice in the example above that if we modify the number of shards, we get different number of samples per GPU and workers even with the implementation of @radulescupetru. This will cause a deadlock in the DDP. So I guess HF expects all shards to contain the same number of samples. Is that a correct assumption @lhoestq?

Setting the number of shards for the datasets above to 2, 2 and 3. Using the `interleave_datasets` I get the following:
```
Rank 0
Assigning 1 shard (or data source) of the dataset to each node.
Dataloader with 0 workers.
Assigning 1 shard (or data source) of the dataset to each node.
775 processed samples
Rank 1
Dataloader with 0 workers.
675 processed samples


Rank 0
Assigning 1 shard (or data source) of the dataset to each node.
Dataloader with 1 workers.
Assigning 1 shard (or data source) of the dataset to each node.
775 processed samples
Rank 1
Dataloader with 1 workers.
675 processed samples


Rank 0
Assigning 1 shard (or data source) of the dataset to each node.
Dataloader with 2 workers.
Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
WARNING:datasets.iterable_dataset:Too many dataloader workers: 2 (max is dataset.num_shards=1). Stopping 1 dataloader workers.
Assigning 1 shard (or data source) of the dataset to each node.
775 processed samples
Rank 1
Dataloader with 2 workers.
675 processed samples


Rank 0
Assigning 1 shard (or data source) of the dataset to each node.
Dataloader with 3 workers.
Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
WARNING:datasets.iterable_dataset:Too many dataloader workers: 3 (max is dataset.num_shards=1). Stopping 2 dataloader workers.
Assigning 1 shard (or data source) of the dataset to each node.
775 processed samples
Rank 1
Dataloader with 3 workers.
675 processed samples


Rank 0
Assigning 1 shard (or data source) of the dataset to each node.
Dataloader with 4 workers.
Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
WARNING:datasets.iterable_dataset:Too many dataloader workers: 4 (max is dataset.num_shards=1). Stopping 3 dataloader workers.
Assigning 1 shard (or data source) of the dataset to each node.
775 processed samples
Rank 1
Dataloader with 4 workers.
675 processed samples
```

I see @LTMeyer, that makes sense. Do you think we should sum the shards by default for concatenating then ? I feel like your use case is more important than ensuring each worker has data of every subdataset in order.

(I wouldn't touch the interleaving logic though)

> Also, I notice in the example above that if we modify the number of shards, we get different number of samples per GPU and workers even with the implementation of @radulescupetru. This will cause a deadlock in the DDP. So I guess HF expects all shards to contain the same number of samples. Is that a correct assumption @lhoestq?

Shards rarely have the same number of samples, so the DDP algorithm itself should be able to stop on its own or have a strategy to circumvent this. For example it can loop until all the nodes have exhausted their data:

```python
def loop():
    while True:
        yield from dataloader
        yield "end"

for x in loop():
    if x == "end":
        exhausted[rank] = True
        continue
    # stop once the data from all the ranks are exhausted
    dist.all_reduce(exhausted)
    if torch.all(exhausted):
        break
    # do your forward pass + loss here
    # model.forward(...)
```

I made a full example here: https://github.com/huggingface/datasets/issues/6623#issuecomment-2379458138

To summarize, and highlight the distinction with https://github.com/huggingface/datasets/pull/7786, there are actually two feature requests:
1. Similarly to `interleave_datasets`, we want to interleave the longest dataset without repetition. This is handled by https://github.com/huggingface/datasets/pull/7786, and is consistant with the rest of the HF features (i.e. `concatenate_datasets` and `interleave_datasets`);
2. We want to be able to _fuse_ datasets and distribute their shards across workers to maximize data loading speed.

 > I feel like your use case is more important than ensuring each worker has data of every subdataset in order.

Indeed my use case, pointed as 2. above is first about maximizing data loading speed and second about mixing the data. The order of priority seems to be the opposite in 1.

> Do you think we should sum the shards by default for concatenating then?

I think the library should at least provide a method for this. Users can then decide what matters the most for their use case (data order or dataloading speed). What do you think?

> Shards rarely have the same number of samples, so the DDP algorithm itself should be able to stop on its own or have a strategy to circumvent this.

If imbalanced data stream in a DDP context is not the responsibility of the datasets library, it is, for me, a reason more to provides a fuse or mix dataset method that sum the shards.

> I made a full example here: https://github.com/huggingface/datasets/issues/6623#issuecomment-2379458138 

Thank you for the example. Pytorch now provides also utilities to handle this problematic case, see [Join context manager in DDP](https://docs.pytorch.org/tutorials/advanced/generic_join.html#:%7E:text=The%20context%20manager%20allows%20the,shadowed%20are%20specified%20by%20hooks)

I'm closing this issue because of several existing solutions:
- https://github.com/huggingface/datasets/pull/7786 allows to interleave datasets without replacement.
- Using [`.shard`](https://huggingface.co/docs/datasets/v4.2.0/en/package_reference/main_classes#datasets.IterableDataset.shard) instead of [`split_dataset_by_node`](https://huggingface.co/docs/datasets/v4.2.0/en/package_reference/main_classes#datasets.distributed.split_dataset_by_node). Given _m_ shards and _n_ ranks, if m % n != 0, the later function will make each of the _n_ ranks go through all of the _m_ shards, although not fetching the same data. On the other hand, the former function can distribute the _m_ shards across the _n_ ranks and make better use of parallel reads.

Thank you @lhoestq and @radulescupetru for the help.