[Trainer] Optimize the use of datasets.IterableDataset in distributed setup
Right now the Trainer uses `IterableDatasetShard` to skip examples on each node and avoid ending up with duplicate data.
This is not efficient for vision or audio tasks since we waste I/O and CPU time reading and decoding files that are not used.

We consider implementing an optimized sharding for distributed training directly in `datasets`. Right now a `datasets.IterableDataset` is already a `torch.utils.data.IterableDataset` that automatically takes care of distributing the necessary input shards to subprocesses in single node (since `datasets` 2.3.0).

The idea would be to also take into account the rank and world size to distribute the input shards. Maybe distributing the. `datasets.IterableDataset` across nodes should be asked explicitly by the user.

cc @sgugger WDYT ? Do you have other ideas in mind to optimize the use of `datasets.IterableDataset` for distributed training in pytorch ?
It seems that `webdataset` also does this when you use `wds.split_by_node` and `wds.split_by_worker`, and is used to train clip or diffusion models at scale.<|||||>That would be awesome! I don't think I need more than an API to tell the iterable dataset that it should take care of the data on process i over n.

Maybe one thing that might be needed is the total length (when available) in some kind of attribute, since the length of the itreable dataset would then be smaller (and the total length is not going to be just a round multiple of the number of processes).<|||||>Opened a PR here, introducing `datasets.distributed.split_dataset_by_node`: https://github.com/huggingface/datasets/pull/5369

feel free to play with it and share your feedbacks :)<|||||>That'll be for when I'm back from vacation ;-)<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.<|||||>Re-ping me in one month GitHub bot ;-)<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.