IterableDataset sharding logic needs improvement
Opened by rwightman (open)

### Describe the bug

The sharding of IterableDatasets with respect to distributed and dataloader worker processes appears problematic with significant performance traps and inconsistencies wrt to distributed train processes vs worker processes.

Splitting across num_workers (per train process loader processes) and world_size (distributed training processes) appears inconsistent.
* worker split: https://github.com/huggingface/datasets/blob/9d6d16117a30ba345b0236407975f701c5b288d4/src/datasets/iterable_dataset.py#L1266-L1283
* distributed split: https://github.com/huggingface/datasets/blob/9d6d16117a30ba345b0236407975f701c5b288d4/src/datasets/iterable_dataset.py#L1335-L1356

In the case of the distributed split, there is a modulus check that flips between two very different behaviours, why is this different than splitting across the data loader workers? For IterableDatasets the DataLoaders worker processes are independent, so whether it's workers within one train process or across a distributed world the shards should be distributed the same, across `world_size * num_worker` independent workers in either case... 

Further, the fallback case when the `n_shards % world_size == 0` check fails is a rather extreme change. I argue it is not desirable to do that implicitly, it should be an explicit case for specific scenarios (ie reliable validation). A train scenario would likely be much better handled with improved wrapping / stopping behaviour to eg also fix #6437. Changing from stepping shards to stepping samples means that every single process reads ALL of the shards. This was never an intended default for sharded training, shards gain their performance advantage in large scale distributed training by explicitly avoiding the need to have every process overlapping in the data they read, by default, only the data allocated to each process via their assigned shards should be read in each pass of the dataset.

Using a large scale CLIP example, some of the larger datasets have 10-20k shards across 100+TB of data. Training with 1000 GPUs we are switching between reading 100 terabytes per epoch to 100 petabytes if say change 20k % 1000 and drop one gpu-node to 20k % 992.

The 'step over samples' case might be worth the overhead in specific validation scenarios where gaurantees of at least/most once samples seen are more important and do not make up a significant portion of train time or are done in smaller world sizes outside of train.

### Steps to reproduce the bug

N/A

### Expected behavior

We have an iterable dataset with N shards, to split across workers
* shuffle shards (same seed across all train processes)
* step shard iterator across distributed processes
* step shard iterator across dataloader worker processes
* shuffle samples in every worker via shuffle buffer (different seed in each worker, but ideally controllable (based on base seed + worker id + epoch).
* end up with (possibly uneven) number of shards per worker but each shard only ever accessed by 1 worker per pass (epoch)



### Environment info

N/A

Comments:
I do not know is it the same probelm as mine.   I think the num_workers should a value of process number for one dataloader mapped to one card, or the total number of  processes for all multiple cards.   
but when I set the num_workers larger then the count of training split files,  it will report  num_workers > n_shards and kill all workers over. as a result, only n_shards workers left, where `n_shard = total files count /  total cards `   
Is that means the num_workers should be the process number on one card? ok,  I changed the num_workers lower, to view it as the number of loader process for one card,  but this time, the data loading is still very slow, it seems that only num_workers  dataloader process are working, not the  num_workers * n_cards  as I thought.   
So how to set a good parameter to make good dataloading?     

Hi @rwightman @carlthome  @NiklasTR @Becomebright @JohnHerry  would it make sense to improve `split_dataset_by_node()` with these changes ?

1) make each node have the same number of samples - this way you don't need barrier / join context
2) ensure each node only loads its own shards (no duplicate, no overlap, no skip)

proposal:
```
    Split a dataset for the node at rank `rank` in a pool of nodes of size `world_size`, with each
    rank having the same number of examples thanks to the `stopping_strategy`.

    The stopping strategy allows each node to have the same number of examples:

    - "first_exhausted": stop when the first node runs of of data, and discard the extra data in the other nodes
    - "all_exhausted": stop when the last node runs out of data, and other nodes may reuse their data to compensate

    For map-style datasets:

    Each node is assigned a chunk of data, e.g. rank 0 is given the first chunk of the dataset.
    To maximize data loading throughput, chunks are made of contiguous data on disk if possible.
    This doesn't need communication between nodes, since each node knows how many examples
    are available and can discard or reuse up to one example accordingly.

    For iterable datasets:

    The shards are evenly assigned across the nodes.
    To maximize data loading throughput, each nodes has its own data and there is no overlap between nodes.
    The stopping strategy has less impact at the end of training if the dataset has a number of shards that is
    a factor of `world_size` (e.g. if `dataset.num_shards % world_size == 0`), since each node has roughly
    the same amount of data available. Nodes communicate using torch distributed to decide when to stop.
```

I opened https://github.com/huggingface/datasets/pull/7854 for those who would like to try this out

So the stopping strategy, there have been versions of this proposed before, I think torch.data had something like this in mind, but it can get complicated, I don't see how you avoid having IPC across dataloaders to achieve something like that robustly. All node/workers need to sync stopping/exhaustion state.

MAP is overall pretty easy because you have a clear view of samples/indices. But you may or may not have a view of locality to make access optimal, or try to balance an approximation of full shuffle with locality of data (if there is any, a folder of jpgs doesn't have any, but samples within a parquet file accessed via map would). 

Iterable, I don't see why shards being a factor of world size makes things easier. You cannot assume that the samples are evenly distributed amongst the shards in any case. And for many sharded datasets you don't have full visiblity on the samples in each shard. You may have an index/map that says N samples in shard X. But when reading you also have to operate reasonably if some of the samples within that shard can't be read or are skipped etc due to an on-the-fly filter. 

To handle those iterable cases gracefully you either need a stopping strategy (global synchronization), or you need a simple solution where you iterate to a target # of steps and ensure the underlying iterator wraps (some repeat) to allow that to happen consistently across all participants. The simple case does not allow one to gaurantee any sort of at-most-once or at-least-once in terms of samples seen.

Another thing here, there's the decision whether to partition samples to nodes once or per epoch (pass over dataset). To be fully shuffled each iteration should result in a deterministic partition of the samples/shards across the nodes. I feel this should be the default. For iterable datasets, each iteration you assign a subset of shards to each node & worker, and then rely on shuffle buffers for sample level shuffle as you iterate over the shards. That is typical TFDS and WDS behaviour. Datasets like MoscaicML's streaming lib sometimes choose to constrain a subset of shards to each node for all of training. It allows an optimization when streaming to cache a subset and then it remains through all of training, but when doing multi-epoch runs it does change the realized shuffle rather significantly. That should be an optimization option not a mandate.

> To handle those iterable cases gracefully you either need a stopping strategy (global synchronization), or you need a simple solution where you iterate to a target # of steps and ensure the underlying iterator wraps (some repeat) to allow that to happen consistently across all participants. The simple case does not allow one to gaurantee any sort of at-most-once or at-least-once in terms of samples seen.

I see. In https://github.com/huggingface/datasets/pull/7854 I opted for global synchronization using torch.distributed. Wdyt ?

> Another thing here, there's the decision whether to partition samples to nodes once or per epoch (pass over dataset). To be fully shuffled each iteration should result in a deterministic partition of the samples/shards across the nodes. I feel this should be the default. For iterable datasets, each iteration you assign a subset of shards to each node & worker, and then rely on shuffle buffers for sample level shuffle as you iterate over the shards. That is typical TFDS and WDS behaviour. Datasets like MoscaicML's streaming lib sometimes choose to constrain a subset of shards to each node for all of training. It allows an optimization when streaming to cache a subset and then it remains through all of training, but when doing multi-epoch runs it does change the realized shuffle rather significantly. That should be an optimization option not a mandate.

On a shuffled dataset, `dataset.set_epoch()` re-shuffles the shards list and re-distributes it among nodes. The seed used for this depends on `shuffle(seed=)` and the epoch number

@lhoestq if you've got global synchronization working well then that's great. Probably some gotchas there?

For IterableDataset Dataloader behaviour in torch, each dataloader worker process runs independently, so each worker is usually given a subset of the shards not just each node.

If you're piggybacking on the training world, not sure that can be done from dataloader workers without possible issues? I would have thought you'd need a separate IPC that's isolated from the main distributed group EDIT: and would need to be initialized from within the worker process via a lazy mechanism (I think)?

Hmm, I guess the workers push their output onto a queue that syncs to the main process via the dataloader iterable, so I guess when you exhaust that the IPC at the train process level probably fine, right. Also means you don't need to worry about balance at the individual worker level... 


I also discovered than Join() lets nodes shutdown when they run out of data while the others continue in DDP. I feel like this would be a general better approach ?

```python
with Join([model]):
    for x in dataloader:
        # training here
```