GitHub issue #5123: datasets freezes with streaming mode in multiple-gpu
State: open | Created: 2022-10-17T03:28:16Z

## Describe the bug
Hi. I am using this dataloader, which is for processing large datasets in streaming mode mentioned in one of examples of huggingface. I am using it to read c4: https://github.com/huggingface/transformers/blob/b48ac1a094e572d6076b46a9e4ed3e0ebe978afc/examples/research_projects/codeparrot/scripts/codeparrot_training.py#L22
During using multi-gpu in accelerator in one node, the code freezes, but works for 1 GPU:   

```
10/16/2022 14:18:46 - INFO - datasets.info - Loading Dataset Infos from /home/jack/.cache/huggingface/modules/datasets_modules/datasets/c4/df532b158939272d032cc63ef19cd5b83e9b4d00c922b833e4cb18b2e9869b01
Steps:   0%|                                                                                                                                                                                                 | 0/400000 [00:00<?, ?it/s]10/16/2022 14:18:47 - INFO - torch.utils.data.dataloader - Shared seed (135290893754684706) sent to store on rank 0
```

# Code to reproduce 
please run this code with `accelerate launch code.py`
```
from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import load_dataset
from torch.utils.data.dataloader import DataLoader
import torch
from datasets import load_dataset
from transformers import AutoTokenizer
import torch
from accelerate.logging import get_logger
from torch.utils.data import IterableDataset
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe


logger = get_logger(__name__)


class ConstantLengthDataset(IterableDataset):
    """
    Iterable dataset that returns constant length chunks of tokens from stream of text files.
        Args:
            tokenizer (Tokenizer): The processor used for proccessing the data.
            dataset (dataset.Dataset): Dataset with text files.
            infinite (bool): If True the iterator is reset after dataset reaches end else stops.
            max_seq_length (int): Length of token sequences to return.
            num_of_sequences (int): Number of token sequences to keep in buffer.
            chars_per_token (int): Number of characters per token used to estimate number of tokens in text buffer.
    """

    def __init__(
        self,
        tokenizer,
        dataset,
        infinite=False,
        max_seq_length=1024,
        num_of_sequences=1024,
        chars_per_token=3.6,
    ):
        self.tokenizer = tokenizer
        # self.concat_token_id = tokenizer.bos_token_id
        self.dataset = dataset
        self.max_seq_length = max_seq_length
        self.epoch = 0
        self.infinite = infinite
        self.current_size = 0
        self.max_buffer_size = max_seq_length * chars_per_token * num_of_sequences
        self.content_field = "text"

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.max_buffer_size:
                    break
                try:
                    buffer.append(next(iterator)[self.content_field])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    if self.infinite:
                        iterator = iter(self.dataset)
                        self.epoch += 1
                        logger.info(f"Dataset epoch: {self.epoch}")
                    else:
                        more_examples = False
                        break
            tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
            all_token_ids = []
            for tokenized_input in tokenized_inputs:
                all_token_ids.extend(tokenized_input)
            for i in range(0, len(all_token_ids), self.max_seq_length):
                input_ids = all_token_ids[i : i + self.max_seq_length]
                if len(input_ids) == self.max_seq_length:
                    self.current_size += 1
                    yield torch.tensor(input_ids)

    def shuffle(self, buffer_size=1000):
        return ShufflerIterDataPipe(self, buffer_size=buffer_size)


def create_dataloaders(tokenizer, accelerator):
    ds_kwargs = {"streaming": True}
    # In distributed training, the load_dataset function gaurantees that only one process
    # can concurrently download the dataset.
    datasets = load_dataset(
        "c4",
        "en",
        cache_dir="cache_dir",
        **ds_kwargs,
    )
    train_data, valid_data = datasets["train"], datasets["validation"]
    with accelerator.main_process_first():
        train_data = train_data.shuffle(buffer_size=10000, seed=None)
        train_dataset = ConstantLengthDataset(
            tokenizer,
            train_data,
            infinite=True,
            max_seq_length=256,
        )
        valid_dataset = ConstantLengthDataset(
            tokenizer,
            valid_data,
            infinite=False,
            max_seq_length=256,
        )
        train_dataset = train_dataset.shuffle(buffer_size=10000)
    train_dataloader = DataLoader(train_dataset, batch_size=160, shuffle=True)
    eval_dataloader = DataLoader(valid_dataset, batch_size=160)
    return train_dataloader, eval_dataloader


def main():
    # Accelerator.
    logging_dir = "data_save_dir/log"
    accelerator = Accelerator(
        gradient_accumulation_steps=1,
        mixed_precision="bf16",
        log_with="tensorboard",
        logging_dir=logging_dir,
    )
    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        accelerator.init_trackers("test")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

    # Load datasets and create dataloaders.
    train_dataloader, _ = create_dataloaders(tokenizer, accelerator)

    train_dataloader = accelerator.prepare(train_dataloader)
    for step, batch in enumerate(train_dataloader, start=1):
        print(step)
    accelerator.end_training()


if __name__ == "__main__":
    main()
```

## Results expected
Being able to run the code for streamining datasets with multi-gpu

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.5.2
- Platform: linux
- Python version: 3.9.12
- PyArrow version: 9.0.0

@lhoestq  I do not have any idea why this freezing happens, and I removed the streaming mode and this was working fine, so I know this is caused by streaming mode of the dataloader part not working well with multi-gpu setting. Since datasets are large, I hope to keep the streamining mode. I very much appreciate your help.