IterableDataset returns duplicated data using PyTorch DDP
Opened by lhoestq (closed)

As mentioned in https://github.com/huggingface/datasets/issues/3423, when using PyTorch DDP the dataset ends up with duplicated data. We already check for the PyTorch `worker_info` for single node, but we should also check for `torch.distributed.get_world_size()` and `torch.distributed.get_rank()`

Comments:
If you use huggingface trainer, you will find the trainer has wrapped a `IterableDatasetShard` to avoid duplication.
See:
https://github.com/huggingface/transformers/blob/dfd818420dcbad68e05a502495cf666d338b2bfb/src/transformers/trainer.py#L835


If you want to support it by datasets natively, maybe we also need to change the code in `transformers` ?

Opened https://github.com/huggingface/transformers/issues/20770 to discuss this :)

Maybe something like this then ?
```python
from datasets.distributed import split_dataset_by_node
ds = split_dataset_by_node(ds, rank=rank, world_size=world_size)
```

For map-style datasets the implementation is trivial (it can simply use `.shard()`).

For iterable datasets we would need to implement a new ExamplesIterable that would only iterate on a subset of the (possibly shuffled and re-shuffled after each epoch) list of shards, based on the rank and world size.

My plan is to skip examples by default to not end up with duplicates.

And if a dataset has a number of shards that is a factor of the world size, then I'd make it more optimized by distributing the shards evenly across nodes instead.

Opened a PR here: https://github.com/huggingface/datasets/pull/5369

feel free to play with it and share your feedbacks :)

@lhoestq I add shuffle after split_dataset_by_node, duplicated data still exist. 
For example, we have a directory named `mock_pretraining_data`, which has three files, `part-00000`, `part-00002`,`part-00002`. 
Text in `part-00000` is like this: 
{"id": 0}
{"id": 1}
{"id": 2}
{"id": 3}
{"id": 4}
{"id": 5}
{"id": 6}
{"id": 7}
{"id": 8}
{"id": 9}

and `part-00001`
{"id": 10}
{"id": 11}
{"id": 12}
{"id": 13}
{"id": 14}
{"id": 15}
{"id": 16}
{"id": 17}
{"id": 18}
{"id": 19}

and `part-00002`
{"id": 20}
{"id": 21}
{"id": 22}
{"id": 23}
{"id": 24}
{"id": 25}
{"id": 26}
{"id": 27}
{"id": 28}
{"id": 29}

And code in `test_dist.py` like this,
```python
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
import os
from transformers import AutoTokenizer, NezhaForPreTraining
from transformers import AdamW, get_linear_schedule_with_warmup
import torch.nn.functional as F
import torch.nn as nn
import torch.distributed as dist
from datasets.distributed import split_dataset_by_node
from torch.nn.parallel import DistributedDataParallel as DDP

os.environ["CUDA_VISIBLE_DEVICES"] = '5,6,7'

dist.init_process_group("nccl")
local_rank = int(os.environ['LOCAL_RANK'])
world_size = torch.distributed.get_world_size()
device = torch.device('cuda', local_rank)
data_dir = './'

def load_trainset(train_path):
    dataset = load_dataset('json', data_dir=os.path.join(data_dir, train_path), split='train', streaming=True)
    return dataset

def collate_fn(examples):
    input_ids = []
    for example in examples:
        input_ids.append(example['id'])
    return torch.LongTensor(input_ids).to(device)


dataset = load_trainset('mock_pretraining_data')
dataset = split_dataset_by_node(dataset, rank=local_rank, world_size=world_size).shuffle(buffer_size=512)
# train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
batch_size = 3
print('batch_size: {}'.format(batch_size))
train_dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)

for x in train_dataloader:
    print({'rank': local_rank, 'id': x})
```
run  `python -m torch.distributed.launch --nproc_per_node=3 test_dist.py`
The output is
```
{'rank': 1, 'id': tensor([12, 15, 14], device='cuda:1')}
{'rank': 1, 'id': tensor([16, 10, 18], device='cuda:1')}
{'rank': 1, 'id': tensor([17, 13, 19], device='cuda:1')}
{'rank': 1, 'id': tensor([11], device='cuda:1')}
{'rank': 0, 'id': tensor([0, 2, 9], device='cuda:0')}
{'rank': 0, 'id': tensor([4, 8, 1], device='cuda:0')}
{'rank': 0, 'id': tensor([5, 3, 6], device='cuda:0')}
{'rank': 0, 'id': tensor([7], device='cuda:0')}
{'rank': 2, 'id': tensor([13, 15, 14], device='cuda:2')}
{'rank': 2, 'id': tensor([19, 17, 18], device='cuda:2')}
{'rank': 2, 'id': tensor([12, 16, 11], device='cuda:2')}
{'rank': 2, 'id': tensor([10], device='cuda:2')}
```
`part-00001` is loaded twice, `part-00002` isn't loaded.

If I run `python -m torch.distributed.launch --nproc_per_node=2 test_dist.py`
The output is weirder，many numbers appear twice
```
{'rank': 1, 'id': tensor([26,  8, 13], device='cuda:1')}
{'rank': 1, 'id': tensor([22, 19, 20], device='cuda:1')}
{'rank': 1, 'id': tensor([12, 28, 11], device='cuda:1')}
{'rank': 1, 'id': tensor([24,  2, 14], device='cuda:1')}
{'rank': 1, 'id': tensor([ 6, 27,  3], device='cuda:1')}
{'rank': 0, 'id': tensor([ 8, 25,  1], device='cuda:0')}
{'rank': 0, 'id': tensor([20,  4, 12], device='cuda:0')}
{'rank': 0, 'id': tensor([14, 29,  5], device='cuda:0')}
{'rank': 0, 'id': tensor([ 7, 18, 23], device='cuda:0')}
{'rank': 0, 'id': tensor([19, 17, 11], device='cuda:0')}
``` 

Hi ! Thanks for reporting, you need to pass `seed=` to `shuffle()` or the processes won't use the same seed to shuffle the shards order before assigning each shard to a node.

The issue is that the workers are not using the same seed to shuffle the shards before splitting the shards list by node.

Opened https://github.com/huggingface/datasets/issues/5696

I have the same issue
```
ds['train'] = load_dataset(streaming=True)
ds['train'] = split_dataset_by_node(ds['train'], rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))
vectorized_datasets = ds.map(
    prepare_dataset,
    remove_columns=raw_datasets_features,
).with_format("torch")

vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(
    buffer_size=500,
    seed=42,
)

def prepare_dataset(batch):
  ....
  print(f"sentence: {batch['sentence']}, target_text: {batch['target_text']}")
  return batch
```
When using split_dataset_by_node(), the data being read is indeed different for each GPU ID.

```
trainer = Trainer(
        model=model,
        data_collator=data_collator,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=vectorized_datasets["train"] if training_args.do_train else None,
        eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,
        tokenizer=processor,
        callbacks=[ShuffleCallback()],
    )
...
train_result = trainer.train(resume_from_checkpoint=checkpoint)
```
However, when I execute trainer.train(), the data being read is different from what I expected.
Because I print the batch value in prepare_dataset() , I observe that the data is the same for each GPU ID.

How should I handle this issue?




There are two ways an iterable dataset can be split by node:
1. if the number of shards is a factor of number of GPUs: in that case the shards are evenly distributed per GPU
2. otherwise, each GPU iterate on the data and at the end keeps 1 sample out of n(GPUs) - skipping the others.

In case 2. it's therefore possible to have the same examples passed to `prepare_dataset` for each GPU.

This doesn't sound optimized though, because it runs the preprocessing on samples that won't be used in the end.

Could you open a new issue so that we can discuss about this and find a solution ?