streaming datasets doesn't work properly with multi-node

[rohitgr7]

Feature request

Let’s say I have a dataset with 5 samples with values [1, 2, 3, 4, 5], with 2 GPUs (for DDP) and batch size of 2. This dataset is an IterableDataset since I am streaming it.

Now I split the dataset using split_dataset_by_node to ensure it doesn’t get repeated. And since it’s already splitted, I don’t have to use DistributedSampler (also they don't work with iterable datasets anyway)?

But in this case I noticed that the:

First iteraton:
first GPU will get → [1, 2]
first GPU will get → [3, 4]

Second iteraton:
first GPU will get → [5]
first GPU will get → Nothing

which actually creates an issue since in case of DistributedSampler, the samples are repeated internally to ensure non of the GPUs at any iteration is missing any data for gradient sync.

So my questions are:

1. Here since splitting is happening before hand, how to make sure each GPU get’s a batch at each iteration to avoid gradient sync issues?
2. Do we need to use DistributedSampler? If yes, how?
3. in the docstrings of split_dataset_by_node, this is mentioned: "If the dataset has a number of shards that is a factor of world_size (i.e. if dataset.n_shards % world_size == 0), then the shards are evenly assigned across the nodes, which is the most optimized. Otherwise, each node keeps 1 example out of world_size, skipping the other examples." Can you explain the last part here?
4. If dataset.n_shards % world_size != 0, is it possible to shard the streaming dataset on the fly to avoid the case where data is missing?

Motivation

Somehow streaming datasets should work with DDP since for big LLMs a lot of data is required and DDP/multi-node is mostly used to train such models and streaming can actually help solve the data part of it.

Your contribution

Yes, I can help in submitting the PR once we get mutual understanding on how it should behave.

[rohitgr7]

@mariosasko, @lhoestq, @albertvillanova
hey guys! can anyone help? or can you guys suggest who can help with this?

[lhoestq]

Hi !

1. When the dataset is running of of examples, the last batches received by the GPU can be incomplete or empty/missing. We haven't implemented yet a way to ignore the last batch. It might require the datasets to provide the number of examples per shard though, so that we can know when to stop.
2. Samplers are not compatible with IterableDatasets in pytorch
3. if dataset.n_shards % world_size != 0 then all the nodes will read/stream the full dataset in order (possibly reading/streaming the same data multiple times), BUT will only yield one example out of world_size so that each example goes to one exactly one GPU.
4. no, sharding should be down up-front and can take some time depending on the dataset size and format

[rohitgr7]

if dataset.n_shards % world_size != 0 then all the nodes will read/stream the full dataset in order (possibly reading/streaming the same data multiple times), BUT will only yield one example out of world_size so that each example goes to one exactly one GPU.

considering there's just 1 shard and 2 worker nodes, do you mean each worker node will load the whole dataset but still receive half of that shard while streaming?

[lhoestq]

Yes both nodes will stream from the 1 shard, but each node will skip half of the examples. This way in total each example is seen once and exactly once during you distributed training.

Though it terms of I/O, the dataset is effectively read/streamed twice.

[rohitgr7]

what if the number of samples in that shard % num_nodes != 0? it will break/get stuck? or is the data repeated in that case for gradient sync?

[lhoestq]

In the case one at least one of the nodes will get an empty/incomplete batch. The data is not repeated in that case. If the training loop doesn't take this into account it can lead to unexpected behaviors indeed.

In the future we'd like to add a feature that would allow the nodes to ignore the last batch, this way all the nodes would only have full batches.

[kkkjyu]

In the case one at least one of the noes will get an empty/incomplete batch. The data is not repeated in that case. If the training loop doesn't take this into account it can lead to unexpected behaviors indeed.

In the future we'd like to add a feature that would allow the nodes to ignore the last batch, this way all the nodes would only have full batches.

Is there any method to modify one dataset's n_shard? modify the number of files is ok? one file == one shard?

[lhoestq]

modify the number of files is ok? one file == one shard?

Yep, one file == one shard :)

[alex-hh]

Hi @lhoestq, do you have any advice on how to implement a fix for the case dataset.n_shards % world_size != 0 while such a fix is not supported in the library?

It seems essential for performing validation in a ddp setting

Simply limiting the number of files is a bit brittle as it relies on world size being consistent to ensure different runs see the same data

How should a user either ignore the last batch or handle the empty batch?

Is the issue of overhanging batches also relevant for map-style datasets?

[lhoestq]

How should a user either ignore the last batch or handle the empty batch?

Check the batch size in the training loop and use all_reduce (or any communication method) to make sure all the nodes got their data before passing them to the model. If some data are missing you can decide to stop the training loop or repeat examples until all the nodes have exhausted their data.

Cc @andrewkho in case you know a way to make the DataLoader stop or add extra samples automatically in case of distributed + unevenly divisible iterable dataset

Is the issue of overhanging batches also relevant for map-style datasets?

The DistributedSampler drops the last data by default to make the dataset evenly divisible.

[andrewkho]

@lhoestq Unfortunately for IterableDataset there isn't a way to do this in general without introducing communciation between ranks, or having all the ranks read all the data before starting to figure out when to stop (which is pretty impractical). My recommendation for these situations where you don't know the total number of samples apriori is to, configure the iterable dataset to yield a fixed number of samples before raising StopIteration, and if necessary, repeat/reshuffle samples to hit that number

[andrewkho]

A heads up that we're planning to land something new in torchdata by end-of-year to help with these scenarios, we'll update this thread when we hvae some code landed

[lhoestq]

I made a quick example with communication between ranks to stop once all the data from all the ranks are exhausted (and repeating data if necessary to end up with a number of samples evenly divisible)

import torch
import torch.distributed as dist
from datasets import Dataset
from datasets.distributed import split_dataset_by_node
from torch.utils.data import DataLoader


# simulate a streaming dataset
num_shards = 1  # change here if you want to simulate a dataset made of many files/shards
ds = Dataset.from_dict({"x": [1, 2, 3, 4, 5]}).to_iterable_dataset(num_shards=num_shards)

# split the dataset for distributed training
dist.init_process_group()
rank, world_size = dist.get_rank(), dist.get_world_size()
ds = split_dataset_by_node(ds, rank=rank,world_size=world_size)
dl = DataLoader(ds)

exhausted = torch.zeros(world_size, dtype=torch.bool)

# IMPORTANT: define the loop to iterate over the local dataset until the data from each rank is exhausted and then return "end"

def loop():
    while True:
        yield from dl
        yield "end"

for x in loop():
    if x == "end":
        exhausted[rank] = True
        continue
    # stop once the data from all the ranks are exhausted
    dist.all_reduce(exhausted)
    if torch.all(exhausted):
        break
    # do your forward pass + loss here
    # model.forward(...)
    print(x)

on my laptop I run torchrun --nnodes=1 --nproc-per-node=2 main.py and I get

{'x': tensor([2])}
{'x': tensor([1])}
{'x': tensor([3])}
{'x': tensor([4])}
{'x': tensor([5])}
{'x': tensor([2])}

we indeed end up with 6 samples, {'x': tensor([2])} was repeated to get 6 examples in total which is divisible by the world size 2.

I also tried with more ranks and with num_workers in DataLoader and it works as expected (don't forget to add if __name__ == '__main__': if necessary for DataLoader multiprocessing)

EDIT: replaced cycle(chain(dl, ["end"])) by loop() after comment #6623 (comment) by @ragavsachdeva

[alex-hh]

great thanks for the example, will give it a try!