Title: Multiprocessing with `Dataset.map` and  `PyTorch` results in deadlock 
State: closed
Labels: description: Something isn't working
name: bug

## Describe the bug
There seems to be an issue with using multiprocessing with `datasets.Dataset.map` (i.e. setting `num_proc` to a value greater than one) combined with a function that uses `torch` under the hood.  The subprocesses that `datasets.Dataset.map` spawns [a this step](https://github.com/huggingface/datasets/blob/1b935dab9d2f171a8c6294269421fe967eb55e34/src/datasets/arrow_dataset.py#L2663) go into wait mode forever.  

## Steps to reproduce the bug
The below code goes into  deadlock when `NUMBER_OF_PROCESSES` is greater than one.
```python

NUMBER_OF_PROCESSES = 2

from transformers import AutoTokenizer, AutoModel
from datasets import load_dataset

dataset = load_dataset("glue", "mrpc", split="train")
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model.to("cpu")

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]

def generate_embeddings_batched(examples):
    sentences_batch = list(examples['sentence1'])
    encoded_input = tokenizer(
        sentences_batch, padding=True, truncation=True, return_tensors="pt"
    )
    encoded_input = {k: v.to("cpu") for k, v in encoded_input.items()}
    model_output = model(**encoded_input)
    embeddings = cls_pooling(model_output)
    examples['embeddings'] = embeddings.detach().cpu().numpy() # 64, 384
    return examples

embeddings_dataset = dataset.map(
    generate_embeddings_batched, 
    batched=True,
    batch_size=10,
    num_proc=NUMBER_OF_PROCESSES
)
```
While debugging it I've seen that it gets "stuck" when calling `torch.nn.Embedding.forward` but some testing shows that the same happens with other functions from `torch.nn`. 

## Environment info
- Platform: Linux-5.14.0-1052-oem-x86_64-with-glibc2.31
- Python version: 3.9.14
- PyArrow version: 9.0.0
- Pandas version: 1.5.0

Not sure if this is a HF problem, a PyTorch problem or something I'm doing wrong..

Thanks!


Comments:
Hi ! Could it be an Out of Memory issue that could have killed one of the processes ? can you check your memory ?

Hi! I don't think it is a memory issue. I'm monitoring the main and spawn python processes and threads with `htop` and the memory does not peak. Besides, the example I've posted above should not be that demanding in terms of memory, right? (I have 32GB of RAM). 

Indeed it should be fine. I couldn't reproduce the error though - I ran your script on my side and it works fine. What version of pytorch are you using ?

Interesting.. I'm using `torch 1.12.1`

I also tried on colab and it works fine 🤔 
Maybe something is wrong with your installation of pytorch ?

Oh actually I just saw that you're using python 3.9

This could be related to https://github.com/huggingface/datasets/issues/4113

We'll fix that as soon as we can, in the meantime you can try to use use single process, or use an older version of python maybe ?

I tried with python 3.7 and the issue persists. In collab, which also uses 3.7 I don't get the issue, so yes I guess is something on mu side... will post it here if I manage to fix it

Hi! Which version of transformers are you using? I test the code on Colab (so python 3.7) with transformers 4.23.1, torch 1.12.1 and pyarrow 9.0.0 (also 6.x), it worked without stuck.

Hi, I have the same problem in use **datasets.IterableDatasetDict.map()**
my pytorch is 2.0.0a0+gitc263bd4
my python is 3.8.16(default, Jun 12 2023, 17:37:21)
work on aarch64 in 16 node, each node with 4*nVidia-A100-40G
every node have 4 process execute code as ↓

```
from datasets import load_dataset, interleave_datasets, IterableDatasetDict, concatenate_datasets
```
...
```
    model_args.cache_dir = '/home/scx/.cache'
    for dataset_name in data_args.datasets_name:
        train_datasets.append(
            load_dataset(
                dataset_name,
                cache_dir=model_args.cache_dir,
                use_auth_token=True if model_args.use_auth_token else None,
                streaming=data_args.streaming,
                split='train'
            ).select_columns('text')
        )
        valid_datasets.append(
            load_dataset(
                dataset_name,
                cache_dir=model_args.cache_dir,
                use_auth_token=True if model_args.use_auth_token else None,
                streaming=data_args.streaming,
                split='validation'
            ).select_columns('text')
        )
        train_dataset = interleave_datasets(train_datasets,
                                        probabilities=data_args.datasets_probabilities, 
                                        seed=training_args.seed,
                                        stopping_strategy='all_exhausted')
    raw_datasets = IterableDatasetDict({'train': train_dataset, 'validation': valid_dataset})
```
...

```
    tokenized_datasets = None
    with training_args.main_process_first(desc="dataset map tokenization"):
        if not data_args.streaming:
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on dataset",
                remove_columns=column_names,
            )
        else:
            #TODO 20230722
            logger.info('{}: {}'.format(__file__, 'tokenized_datasets = raw_datasets.map('))
            logger.info('len raw_datasets: {}'.format(len(raw_datasets.items())))
            logger.info('raw_datasets:{}'.format(raw_datasets.items()))
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                batch_size=1000,
                remove_columns=column_names
            )
            logger.info('map ok!')
            logger.info('show train: {}'.format(next(iter(tokenized_datasets['train']))))
            logger.info('ok')
            # ### RAW CODE ###
            # tokenized_datasets = raw_datasets.map(
            #     tokenize_function,
            #     batched=True,
            #     batch_size=1000,
            #     remove_columns=column_names
            # )
            #TODO 20230722
    logger.info("Finish tokenization")
```
the output of my code is
```
07/22/2023 21:57:09 - INFO - __main__ - /demo/run_blue_space.py: tokenized_datasets = raw_datasets.map(
07/22/2023 21:57:09 - INFO - __main__ - len raw_datasets: 2
07/22/2023 21:57:09 - INFO - __main__ - raw_datasets:dict_items([('train', <datasets.iterable_dataset.IterableDataset object at 0x4005ee301190>), ('validation', <datasets.iterable_dataset.IterableDataset object at 0x4005ee5427f0>)])
07/22/2023 21:57:09 - INFO - __main__ - map ok!
07/22/2023 22:01:07 - INFO - __main__ - show train: {'input_ids': [14608, 26797, 31891, 34260, 12227, 33207, 5, 5, 31632, 26797, 31891, 34260, 12227, 33207, 7398, 28561, 31236, 31177, 31253, 33558, 31556, 31377, 72, 20732, 32383, 32295, 14027, 31178, 53, 61, 53, 55, 31189, 31146, 31321, 31235, 53, 61, 56, 58, 31189, 31145, 72, 53, 61, 58, 54, 31189, 54, 31245, 53, 60, 31224, 31896, 31178, 28561, 29331, 20732, 31888, 32637, 4426, 2824, 72, 53, 61, 60, 55, 31189, 53, 54, 31245, 53, 31224, 31896, 31178, 28561, 29331, 26137, 20732, 4426, 2824, 73, 54, 52, 52, 52, 31189, 61, 31245, 59, 31224, 31896, 31178, 29331, 28561, 20732, 4426, 2824, 73, 5], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
07/22/2023 22:01:07 - INFO - __main__ - ok
```



@bio-punk `IterableDatasetDict.map` does not support multiprocessing (only `DatasetDict.map` and `Dataset.map` do), so please open a new issue as this doesn't seem to be related to the original issue. 

Closing as this issue doesn't seem to be related to `datasets`.