Concatenate IterableDataset instances and distribute underlying shards in a RoundRobin manner

[LTMeyer]

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.
• 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.

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.

[LTMeyer]

With datasets.Dataset

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

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}

[LTMeyer]

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])}

Experiment script

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}")

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.

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:

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?

[LTMeyer]

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.

[lhoestq]

I believe this PR could solve your issue? :)

#7786

[LTMeyer]

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).

[radulescupetru]

@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]

@LTMeyer With the following script and using the code from #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.

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, 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

[lhoestq]

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.

[LTMeyer]

Can you confirm it works for you now @LTMeyer ?

Result with a547d81:

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 #7792 (comment) 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.

[radulescupetru]

@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]

@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.

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 455bfaa:

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.

[lhoestq]

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.