Struk during training on multi-GPUs with deepspeed, but work with single GPU
Author: loveunk
Pull request: no

### System Info

- `transformers` version: 4.31.0
- Platform: Linux-4.15.0-156-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.20.1
- Safetensors version: 0.4.1
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@pacman100 @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

During the training with the offical LLaVA code (https://github.com/haotian-liu/LLaVA), it stucks with multi-GPUs (2 or 4 A6000), but works with a single GPU.

Here's my script:
```
deepspeed  --include localhost:2,3 --master_port 61000 llava/train/train_mem.py \
    --lora_enable True --lora_r 128 --lora_alpha 256 --mm_projector_lr 2e-5 \
    --deepspeed ./scripts/zero2.json \
    --model_name_or_path /drive/models/llava-v1.5-13b \
    --version v1 \
    --data_path ./playground/data/extract_textvqa.json \
    --image_folder ./playground/data \
    --vision_tower openai/clip-vit-large-patch14-336 \
    --mm_projector_type mlp2x_gelu \
    --mm_vision_select_layer -2 \
    --mm_use_im_start_end False \
    --mm_use_im_patch_token False \
    --image_aspect_ratio pad \
    --group_by_modality_length True \
    --bf16 True \
    --output_dir ./checkpoints/llava-v1.5-13b-task-lora-2gpu \
    --num_train_epochs 1 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 50000 \
    --save_total_limit 1 \
    --learning_rate 2e-4 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --tf32 True \
    --model_max_length 2048 \
    --gradient_checkpointing True \
    --dataloader_num_workers 4 \
    --lazy_preprocess True \
    --report_to wandb
```

zero2.json:
```
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "bf16": {
        "enabled": "auto"
    },
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto"
    }
}
```

It stuck after loading the model, here the call stack with `faulthandler.dump_traceback_later(20, repeat=True)`:

```
Timeout (0:00:20)!
Thread 0x00007f8570ecc700 (most recent call first):
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 324 in wait
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 607 in wait
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/tqdm/_monitor.py", line 60 in run
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 1016 in _bootstrap_inner
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 973 in _bootstrap

Thread 0x00007f87f88404c0 (most recent call first):
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1570 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1451 in wrapper
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/torch.py", line 129 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 216 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 116 in log_wrapper
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1012 in _broadcast_model
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1084 in _configure_distributed_model
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 267 in __init__
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/__init__.py", line 165 in initialize
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/accelerate/accelerator.py", line 1537 in _prepare_deepspeed
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/accelerate/accelerator.py", line 1198 in prepare
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/transformers/trainer.py", line 1656 in _inner_training_loop
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/transformers/trainer.py", line 1539 in train
  File "/drive/LLaVA/llava/train/train.py", line 935 in train
  File "/drive/LLaVA/llava/train/train_mem.py", line 16 in <module>
Timeout (0:00:20)!
Thread 0x00007f5f4cacd700 (most recent call first):
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 324 in wait
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 607 in wait
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/tqdm/_monitor.py", line 60 in run
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 1016 in _bootstrap_inner
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/threading.py", line 973 in _bootstrap

Thread 0x00007f61d0f444c0 (most recent call first):
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1570 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1451 in wrapper
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/torch.py", line 129 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 216 in broadcast
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 116 in log_wrapper
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1012 in _broadcast_model
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1084 in _configure_distributed_model
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 267 in __init__
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/deepspeed/__init__.py", line 165 in initialize
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/accelerate/accelerator.py", line 1537 in _prepare_deepspeed
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/accelerate/accelerator.py", line 1198 in prepare
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/transformers/trainer.py", line 1656 in _inner_training_loop
  File "/home/vscode/miniconda3/envs/llava/lib/python3.10/site-packages/transformers/trainer.py", line 1539 in train
  File "/drive/LLaVA/llava/train/train.py", line 935 in train
  File "/drive/LLaVA/llava/train/train_mem.py", line 16 in <module>
```

### Expected behavior

should do training

Comments:
tried to run the following code, stuck with backend `nccl`, but works with `gloo`:

```python
import torch.distributed as dist
import argparse
import torch

torch.cuda.set_device(int(os.environ['LOCAL_RANK']))
device = torch.device("cuda", int(os.environ['LOCAL_RANK']))

dist.init_process_group("nccl") 
dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
```

run command: 
```shell
python -m torch.distributed.launch --nproc_per_node=2 test.py
```

It looks like there are some NCCL issues with my GPU server.
The training works after turning off the GPU P2P communication with `export NCCL_P2P_DISABLE=1`.

FYI: https://blog.csdn.net/qq_40947610/article/details/128118180

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.