Language model example hanging
Author: roywei
Pull request: no
Labels: name: bug

### System Info

```shell
Transformer: 4.17
PyTorch: 1.10.2

Public Docker image I used:
763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:1.10.2-transformers4.17.0-gpu-py38-cu113-ubuntu20.04-v1.0

Hardware:
1 AWS EC2 P4D instance, 8 A100 GPUs
```


### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Just run the examples using command in readme

https://github.com/huggingface/transformers/tree/v4.17.0/examples/pytorch/language-modeling

```
python run_mlm.py \
    --model_name_or_path roberta-base \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --do_train \
    --do_eval \
    --output_dir /tmp/test-mlm
```

It hangs after finish downloading the data.

```
root/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-7f695945617
eacf7.arrow
Grouping texts in chunks of 512: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████
██████████████████████████████████| 4/4 [00:00<00:00, 14.70ba/s]
04/20/2022 18:15:03 - INFO - datasets.utils.file_utils - https://raw.githubusercontent.com/huggingface/datasets/1.18.4/metrics/accuracy/accuracy.py
not found in cache or force_download set to True, downloading to /root/.cache/huggingface/datasets/downloads/tmp0f76eppv
Downloading: 3.19kB [00:00, 3.66MB/s]
```


### Expected behavior

```shell
run training as expected
```


Comments:
Hello @roywei, 

The example scripts are always tested and maintained with the corresponding `transformers` version. 

I couldn't reproduce your error! Here is how I ran the script
```python
import sagemaker
from sagemaker.huggingface import HuggingFace

# gets role for executing training job
role = sagemaker.get_execution_role()
hyperparameters = {
	'model_name_or_path':'roberta-base',
	'output_dir':'/opt/ml/model',
    'dataset_name':'wikitext',
    'dataset_config_name': 'wikitext-2-raw-v1',
    'per_device_train_batch_size': 8,
    'per_device_eval_batch_size': 8,
    'do_train': True,
    'do_eval': True
	# add your remaining hyperparameters
	# more info here https://github.com/huggingface/transformers/tree/v4.17.0/examples/pytorch/language-modeling
}

# git configuration to download our fine-tuning script
git_config = {'repo': 'https://github.com/huggingface/transformers.git','branch': 'v4.17.0'}

# creates Hugging Face estimator
huggingface_estimator = HuggingFace(
	entry_point='run_mlm.py',
	source_dir='./examples/pytorch/language-modeling',
	instance_type='ml.p3.2xlarge',
	instance_count=1,
	role=role,
	git_config=git_config,
	transformers_version='4.17.0',
	pytorch_version='1.10.2',
	py_version='py38',
	hyperparameters = hyperparameters
)

# starting the train job
huggingface_estimator.fit()
```

P.S. Is there a reason why you don't use Amazon SageMaker? 

Hi I can confirm on single node it's working as expected. I'd like to rule out some possibilities on 2 node run, so trying to run it with DDP + NCCL on 2 nodes.

I added the following env var in `run_mlm.py` to help with pytorch  ddp launch

```
os.environ["LOCAL_RANK"] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]
os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"]
os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
os.environ["MASTER_ADDR"]="algo-1"
os.environ["MASTER_PORT"]="12345"
```

also  added distribution in estimator API 

```
distribution = {"mpi":{"enabled":True}}
huggingface_estimator = HuggingFace(
	entry_point='run_mlm.py',
	source_dir='./examples/pytorch/language-modeling',
	instance_type='ml.p3.2xlarge',
	instance_count=1,
	role=role,
	git_config=git_config,
	transformers_version='4.17.0',
	pytorch_version='1.10.2',
	py_version='py38',
        distribution=distribution,
	hyperparameters = hyperparameters
)
```

Above will produce a hang, let me know if I did anything wrong. Thanks!


I think the key difference is if I use the distributed launcher, the job hanged at beginning of the training, if I specify the distribution and num of nodes to 1, it would still hang. 
```
mpirun -N  xxx python run_mlm.py
```
 If I do not specify distribution arg, it will be ran using `python run_mlm.py` directly and it works fine on 1 node.

Hi, I've found the issue, it's from our user's modification of the script, that's setting `preprocessing_num_workers` to a very large number, leading to the hang. Removing that makes both NCCL and SMDDP works. The defaults in this example is working fine as well. 

Thanks for all the help @philschmid .
Closing!

Hey @roywei,

not sure if you knew this but all `examples/` scripts, including `run_mlm.py` support distributed training (data & model parallelism) out of the box through the `Trainer`, here is an example notebook for Data parallelism
* https://github.com/huggingface/notebooks/blob/main/sagemaker/03_distributed_training_data_parallelism/sagemaker-notebook.ipynb

Thanks a lot! I will pass this as a reference to our user.