Multiprocessing with `Dataset.map` and  `PyTorch` results in deadlock 
MEMBER
Hi ! Could it be an Out of Memory issue that could have killed one of the processes ? can you check your memory ?
2022-10-10T15:37:10Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273498202
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6Apa
https://api.github.com/repos/huggingface/datasets/issues/comments/1273498202/reactions
2022-10-10T15:37:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273498202
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
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). 
2022-10-10T15:44:58Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273507488
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6C6g
https://api.github.com/repos/huggingface/datasets/issues/comments/1273507488/reactions
2022-10-10T15:44:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273507488
https://avatars.githubusercontent.com/u/36822895?v=4
https://api.github.com/users/RR-28023/events{/privacy}
https://api.github.com/users/RR-28023/followers
https://api.github.com/users/RR-28023/following{/other_user}
https://api.github.com/users/RR-28023/gists{/gist_id}
https://github.com/RR-28023
RR-28023
MDQ6VXNlcjM2ODIyODk1
https://api.github.com/users/RR-28023/orgs
https://api.github.com/users/RR-28023/received_events
https://api.github.com/users/RR-28023/repos
https://api.github.com/users/RR-28023/starred{/owner}{/repo}
https://api.github.com/users/RR-28023/subscriptions
User
https://api.github.com/users/RR-28023
public
MEMBER
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 ?
2022-10-10T16:14:30Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273544824
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6MB4
https://api.github.com/repos/huggingface/datasets/issues/comments/1273544824/reactions
2022-10-10T16:14:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273544824
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Interesting.. I'm using `torch 1.12.1`
2022-10-10T16:21:03Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273552279
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6N2X
https://api.github.com/repos/huggingface/datasets/issues/comments/1273552279/reactions
2022-10-10T16:21:03Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273552279
https://avatars.githubusercontent.com/u/36822895?v=4
https://api.github.com/users/RR-28023/events{/privacy}
https://api.github.com/users/RR-28023/followers
https://api.github.com/users/RR-28023/following{/other_user}
https://api.github.com/users/RR-28023/gists{/gist_id}
https://github.com/RR-28023
RR-28023
MDQ6VXNlcjM2ODIyODk1
https://api.github.com/users/RR-28023/orgs
https://api.github.com/users/RR-28023/received_events
https://api.github.com/users/RR-28023/repos
https://api.github.com/users/RR-28023/starred{/owner}{/repo}
https://api.github.com/users/RR-28023/subscriptions
User
https://api.github.com/users/RR-28023
public
MEMBER
I also tried on colab and it works fine 🤔 
Maybe something is wrong with your installation of pytorch ?
2022-10-10T16:48:01Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273581467
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6U-b
https://api.github.com/repos/huggingface/datasets/issues/comments/1273581467/reactions
2022-10-10T16:48:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273581467
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
MEMBER
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 ?
2022-10-10T16:49:09Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273582666
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6VRK
https://api.github.com/repos/huggingface/datasets/issues/comments/1273582666/reactions
2022-10-10T16:49:54Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273582666
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
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
2022-10-10T18:31:47Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1273677294
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5L6sXu
https://api.github.com/repos/huggingface/datasets/issues/comments/1273677294/reactions
2022-10-10T18:31:47Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1273677294
https://avatars.githubusercontent.com/u/36822895?v=4
https://api.github.com/users/RR-28023/events{/privacy}
https://api.github.com/users/RR-28023/followers
https://api.github.com/users/RR-28023/following{/other_user}
https://api.github.com/users/RR-28023/gists{/gist_id}
https://github.com/RR-28023
RR-28023
MDQ6VXNlcjM2ODIyODk1
https://api.github.com/users/RR-28023/orgs
https://api.github.com/users/RR-28023/received_events
https://api.github.com/users/RR-28023/repos
https://api.github.com/users/RR-28023/starred{/owner}{/repo}
https://api.github.com/users/RR-28023/subscriptions
User
https://api.github.com/users/RR-28023
public
NONE
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.
2022-10-18T16:18:53Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1282656625
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5Mc8lx
https://api.github.com/repos/huggingface/datasets/issues/comments/1282656625/reactions
2022-10-18T16:18:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1282656625
https://avatars.githubusercontent.com/u/11248271?v=4
https://api.github.com/users/yuancu/events{/privacy}
https://api.github.com/users/yuancu/followers
https://api.github.com/users/yuancu/following{/other_user}
https://api.github.com/users/yuancu/gists{/gist_id}
https://github.com/yuancu
yuancu
MDQ6VXNlcjExMjQ4Mjcx
https://api.github.com/users/yuancu/orgs
https://api.github.com/users/yuancu/received_events
https://api.github.com/users/yuancu/repos
https://api.github.com/users/yuancu/starred{/owner}{/repo}
https://api.github.com/users/yuancu/subscriptions
User
https://api.github.com/users/yuancu
public
NONE
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
```


2023-07-22T14:16:09Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1646593734
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5iJQbG
https://api.github.com/repos/huggingface/datasets/issues/comments/1646593734/reactions
2023-07-22T14:16:09Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1646593734
https://avatars.githubusercontent.com/u/9584268?v=4
https://api.github.com/users/bio-punk/events{/privacy}
https://api.github.com/users/bio-punk/followers
https://api.github.com/users/bio-punk/following{/other_user}
https://api.github.com/users/bio-punk/gists{/gist_id}
https://github.com/bio-punk
bio-punk
MDQ6VXNlcjk1ODQyNjg=
https://api.github.com/users/bio-punk/orgs
https://api.github.com/users/bio-punk/received_events
https://api.github.com/users/bio-punk/repos
https://api.github.com/users/bio-punk/starred{/owner}{/repo}
https://api.github.com/users/bio-punk/subscriptions
User
https://api.github.com/users/bio-punk
public
COLLABORATOR
@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. 
2023-07-24T15:28:30Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1648136442
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5iPJD6
https://api.github.com/repos/huggingface/datasets/issues/comments/1648136442/reactions
2023-07-24T15:28:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1648136442
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
COLLABORATOR
Closing as this issue doesn't seem to be related to `datasets`.
2023-07-24T15:29:13Z
https://github.com/huggingface/datasets/issues/5094#issuecomment-1648137707
https://api.github.com/repos/huggingface/datasets/issues/5094
IC_kwDODunzps5iPJXr
https://api.github.com/repos/huggingface/datasets/issues/comments/1648137707/reactions
2023-07-24T15:29:13Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1648137707
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
## 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!
