Support concatenating multiple streaming datasets while preserving the sum of shards

[muyihao]

Feature request

It would be useful if concatenating multiple IterableDatasets could preserve the shard information from each input dataset, so that:

• the total number of shards equals the sum of the input shards;
• split_dataset_by_node() and worker sharding continue to work efficiently;
• downstream distributed dataloaders can fully exploit the available parallelism.

Motivation

Currently, it's straightforward to concatenate multiple streaming datasets using interleave_datasets or other approaches. However, after concatenation, the resulting dataset does not preserve the original sharding information from each input dataset.

For distributed training, especially with large-scale datasets, this leads to a reduced number of effective shards compared to the sum of the input datasets' shards. Consequently, data loading parallelism may be limited, and the available worker concurrency cannot fully utilize the original dataset partitioning.

For example:

• IterableDataset A: 1 shard
• IterableDataset B: 1 shard

After concatenation, we would expect the resulting streaming dataset to expose 2 shards, allowing distributed workers to shard the combined dataset naturally.

Your contribution

I'd be happy to work on this feature and submit a PR if the maintainers think this would be a useful addition.

Any feedback on the proposed design would be greatly appreciated.

[lhoestq]

it's already the case afaik:

In [1]: from datasets import IterableDataset, concatenate_datasets

In [2]: ds1 = IterableDataset.from_dict({"i": [0, 1]})

In [3]: ds2 = IterableDataset.from_dict({"i": [2, 3]})

In [4]: ds = concatenate_datasets([ds1, ds2])

In [5]: list(ds)
Out[5]: [{'i': 0}, {'i': 1}, {'i': 2}, {'i': 3}]

In [6]: ds.num_shards  # = sum([ds1.num_shards, ds2.num_shards])
Out[6]: 2

On the other hand interleave_datasets() doesn't result in the sum of shards. Instead each new shard is made of at least one shard of every input dataset so that each shard effectively contains interleaved data.

In [7]: from datasets import interleave_datasets

In [8]: ds = interleave_datasets([ds1, ds2], stopping_strategy="all_exhausted")

In [9]: list(ds)
Out[9]: [{'i': 0}, {'i': 2}, {'i': 1}, {'i': 3}]

In [10]: ds.num_shards  # = min([ds1.num_shards, ds2.num_shards])
Out[10]: 1