# Keeping IterableDataset node-wise split fixed during DDP

**URL:** <https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713>\
**Category:** 🤗Datasets\
**Created:** [October 16, 2023, 11:05am UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713 "2023-10-16T11:05:07Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![s9ngularity](https://avatars.discourse-cdn.com/v4/letter/s/90db22/32.png) [@s9ngularity](https://discuss.huggingface.co/u/s9ngularity)\
**Post date:** [October 16, 2023, 11:05am UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/1 "2023-10-16T11:05:07Z")

</div>

Hi all,

I’m using `datasets.IterableDataset` (more specifically `IterableDataset.from_generator`). I have been using it with DDP Pytorch by streaming all the data across multiple GPU nodes from a source as caching the whole dataset locally requires **too much disk space**.

However, if I could shard the dataset per DDP node, then my data could fit on disk. So I am wondering is it possible use `IterableDataset.from_generator` so that:

1. Each DDP node gets assigned a **fixed subset** of the shards (e.g., based on a seed) for the whole run.
2. and each DDP node shuffles its shards for each epoch.

If it is possible let me know how to make it.

Thank you!

---

<div class="post-metadata">

**Author:** ![mariosasko](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mariosasko/32/31548_2.png) [@mariosasko](https://discuss.huggingface.co/u/mariosasko)\
**Post date:** [October 16, 2023, 5:23pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/2 "2023-10-16T17:23:28Z")

</div>

You can use `shuffle` and `set_epoch` to shuffle the shards and samples in between epochs (explained [here](https://huggingface.co/docs/datasets/v2.14.5/en/stream#stream) in the docs) and [`split_dataset_by_node`](https://huggingface.co/docs/datasets/v2.14.5/en/package_reference/main_classes#datasets.distributed.split_dataset_by_node) to split the dataset across nodes.

For this to work efficiently, the dataset must consist of many shards (`n_shards` returns the number of shards; `dataset.n_shards % world_size == 0` is the ideal number). An example of creating a sharded dataset is available [here](https://huggingface.co/docs/datasets/v2.14.5/en/package_reference/main_classes#datasets.Dataset.from_generator.example-2) (shards are formed by sharding `gen_kwargs`’s list values).

---

<div class="post-metadata">

**Author:** ![s9ngularity](https://avatars.discourse-cdn.com/v4/letter/s/90db22/32.png) [@s9ngularity](https://discuss.huggingface.co/u/s9ngularity)\
**Post date:** [October 16, 2023, 5:33pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/3 "2023-10-16T17:33:03Z")

</div>

Thanks, yes, the dataset has many shards.

But could you give the example how to keep the dataset split fixed for each node during the whole DDP run?

The documentation had the following example for `split_dataset_by_node`:

```python
from datasets.distributed import split_dataset_by_node
ids = ds.to_iterable_dataset(num_shards=512)
ids = ids.shuffle(buffer_size=10_000) # will shuffle the shards order and use a shuffle buffer when you start iterating
ids = split_dataset_by_node(ids, world_size=8, rank=0) # will keep only 512 / 8 = 64 shards from the shuffled lists of shards when you start iterating
dataloader = torch.utils.data.DataLoader(ids, num_workers=4) # will assign 64 / 4 = 16 shards from this node's list of shards to each worker when you start iterating
for example in ids:
    pass

```

However, if I understand correctly, in this example each node would get a new random subset of shards every epoch, and then my nodes cannot to re-use the previous local cache.

Thanks again!

---

<div class="post-metadata">

**Author:** ![mariosasko](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mariosasko/32/31548_2.png) [@mariosasko](https://discuss.huggingface.co/u/mariosasko)\
**Post date:** [October 16, 2023, 5:50pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/4 "2023-10-16T17:50:01Z")

</div>

Like this:

```python
from datasets.distributed import split_dataset_by_node
ids = ds.to_iterable_dataset(num_shards=512)
ids = ids.shuffle(buffer_size=10_000) # will shuffle the shards order and use a shuffle buffer when you start iterating
ids = split_dataset_by_node(ids, world_size=8, rank=0) # will keep only 512 / 8 = 64 shards from the shuffled lists of shards when you start iterating
dataloader = torch.utils.data.DataLoader(ids, num_workers=4) # will assign 64 / 4 = 16 shards from this node's list of shards to each worker when you start iterating
for epoch in range(num_epochs):
    ids.set_epoch(epoch) # operates on the dataset split
    for example in dataloader:
        pass

```

---

<div class="post-metadata">

**Author:** ![s9ngularity](https://avatars.discourse-cdn.com/v4/letter/s/90db22/32.png) [@s9ngularity](https://discuss.huggingface.co/u/s9ngularity)\
**Post date:** [October 16, 2023, 6:12pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/5 "2023-10-16T18:12:06Z")

</div>

Unfortunately, this does not work. It gives different subset of shards each epoch for each node.

---

<div class="post-metadata">

**Author:** ![mariosasko](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mariosasko/32/31548_2.png) [@mariosasko](https://discuss.huggingface.co/u/mariosasko)\
**Post date:** [October 18, 2023, 6:00pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/6 "2023-10-18T18:00:15Z")

</div>

I think this is for consistency with PyTorch’s `DistributedSampler.set_epoch` (@lhoestq should know more as the author of this feature).

If the dataset in question is a Hub dataset without a loading script, one option is to manually shuffle and split the files/shards to ensure they stay fixed in each epoch.

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [October 18, 2023, 9:45pm UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/7 "2023-10-18T21:45:57Z")

</div>

> [@s9ngularity](#):
>
> However, if I understand correctly, in this example each node would get a new random subset of shards every epoch, and then my nodes cannot to re-use the previous local cache.

That’s correct: at each epoch the shards are re-shuffled across all the nodes. Assigning certain shards to certain nodes is not supported

---

<div class="post-metadata">

**Author:** ![macwiatrak](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/macwiatrak/32/25095_2.png) [@macwiatrak](https://discuss.huggingface.co/u/macwiatrak)\
**Post date:** [March 22, 2024, 10:02am UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/8 "2024-03-22T10:02:26Z")

</div>

If we are using an IterableDataset, does the HuggingFace trainer automatically split the datasets by node, or is it something one should do manually? See: [How to handle IterableDataset with HuggingFace trainer and num\_workers in DDP setup](https://discuss.huggingface.co/t/how-to-handle-iterabledataset-with-huggingface-trainer-and-num-workers-in-ddp-setup/78459)

---

<div class="post-metadata">

**Author:** ![dpiscia](https://avatars.discourse-cdn.com/v4/letter/d/b5e925/32.png) [@dpiscia](https://discuss.huggingface.co/u/dpiscia)\
**Post date:** [April 29, 2024, 6:17am UTC](https://discuss.huggingface.co/t/keeping-iterabledataset-node-wise-split-fixed-during-ddp/58713/9 "2024-04-29T06:17:59Z")

</div>

Flower dataset [Flower Datasets 0.1.0](https://flower.ai/docs/datasets/) which is built on top of huggingface dataset probably solve that issue.