IterableDataset supports automatic sharding

[howitry]

Feature request

Added sharding function support to the streaming IterableDataset, allowing users to adjust the number of shards according to their training resources. For example:

dataset = load_dataset(*, stream=True)
dataset = dataset.shard(num_shards=num_shards, num_samples=num_samples) #We may need to know the total number of samples (num_samples) in advance.

Motivation

When performing large-scale pre-training in a distributed environment, large datasets may only be loaded in a streaming manner. To improve training efficiency, my current approach is as follows:

file_type="parquet"
dataset_path="./*.parquet"
dataset = load_dataset(file_type,data_files=dataset_path, stream=True)
dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size)

I split a large file into N = world_size * dataloader_num_workers files and placed them under dataset_path. This ensures that each GPU processes different shards. However, this approach has some issues. If the number of GPUs used to train the model changes next time, I need to split the large file again to ensure that IterableDataset.num_shards = world_size * dataloader_num_workers.
I'd like to know if there's a better approach, such as directly loading the large dataset in a streaming manner and then sharding the IterableDataset based on the number of GPUs and num_workers, similar to the approach in Example 1 of https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset @lhoestq

[lhoestq]

You can already use .shard() instead like this:

dataset = dataset.shard(index=rank, num_shards=world_size)

note that it requires that dataset.num_shards >= world_size, and that it may result in nodes having the same number of shards +- 1

[howitry]

You can already use .shard() instead like this:

dataset = dataset.shard(index=rank, num_shards=world_size)
note that it requires that dataset.num_shards >= world_size, and that it may result in nodes having the same number of shards +- 1

This means I have to ensure that the initial num_shards is greater than the number of GPUs I use each time, which seems inflexible. Is there a way to dynamically divide the data into multiple shards based on the number of GPUs used each time? For example:

dataset = load_dataset(*, stream=True) # dataset.num_shards()=1
num_shards=world_size*dataloader_num_workers
dataset = dataset.dynamically_shard(num_shards=num_shards, num_samples=num_samples) #We may need to know the total number of samples (num_samples) in advance.

[lhoestq]

Is there a way to dynamically divide the data into multiple shards based on the number of GPUs used each time?

No it's not possible without either

1. doing data skipping, which degrades the data loading performance significantly (every node has to download the same data and skip most samples)
2. or divide the original files further, which requires additional logic for every file format

I would be interested in exploring 2 though, maybe if we start with Parquet support. Right now it fails because ArrowExamplesIterable doesn't know how to shard more than num_shards. We could have instead a ReshardableArrowExamplesIterable that would pass the right arguments to _generate_tables() in parquet.py to only read the data requested for a specific node

[howitry]

ReshardableArrowExamplesIterable

Okay, my datasets are all on my local disk, so I haven't considered the overhead of data download. Are there any tutorials on creating custom iterable datasets? For example, a custom iterabledataset.__iter__ function can be used to skip data, and it can inherit operations like iterabledataset.map.