FSDP hangs before training starts
### System Info

- `transformers` version: 4.27.4
- Platform: Linux-5.4.0-128-generic-x86_64-with-glibc2.31
- Python version: 3.9.16
- Huggingface_hub version: 0.13.4
- PyTorch version (GPU?): 2.0.0+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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@pacman100 @stas @sgugger

### 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

Steps to reproduce the behaviour : 

1. I created the FSDP Config file using accelerate config as follows : 
```
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: GPTJBlock
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 6
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

```

2. My bash script looks like this : 

```
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5

accelerate launch train_llm.py \
--output_dir /path/to/dir  \
--model_name_or_dir "EleutherAI/gpt-j-6B" \
--do_train --per_device_train_batch_size 8 \
--do_eval --per_device_eval_batch_size 8 \
--num_train_epochs 3 \
--evaluation_strategy "steps" \
--eval_steps 1000 \
--save_strategy "steps" \
--save_steps 1000 \
--learning_rate 5e-5 \
--logging_steps 1 \
--bf16 \
--run_name run_fsdp \
--gradient_checkpointing true \
--warmup_ratio 0.03 \
--fsdp "full_shard auto_wrap" \
--fsdp_transformer_layer_cls_to_wrap "GPTJBlock" 
```

3. My train_llm.py file look like this this  - 

```
if __name__ == "__main__":
    
    parser = HfArgumentParser(TrainingArguments)
    parser.add_argument("--model_name_or_dir")
    
    training_args, args = parser.parse_args_into_dataclasses()

    transformers.logging.set_verbosity_debug()

    model = AutoModelForCausalLM.from_pretrained(args.model_name_or_dir, use_cache=True, ignore_mismatched_sizes=True)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_dir, use_cache=True)
    tokenizer.pad_token_id = tokenizer.eos_token_id

 
    train_path = 'path/to/train'
    train_data = glob(train_path)
    
    val_path = 'path/to/val'
    val_data = glob(val_path)

    dataset = load_dataset("json", data_files = {"train": train_data, "validation" : val_data})
    dataset = dataset.map(transform, batched=True, remove_columns = ["id" ,"tokens"])
    
    train_dataset = dataset["train"]
    val_dataset = dataset["validation"]

    trainer = Trainer(
        model,
        training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorForTokenClassification(tokenizer, padding='longest'),
        compute_metrics=None, 
        callbacks = [TensorBoardCallback()]      
    )

    if trainer.is_world_process_zero():
        print(dataset)
        
    trainer.pop_callback(MLflowCallback)
    
    if training_args.do_train:
        if trainer.is_world_process_zero():
            print("Training...")

        start = time.time()
        trainer.train()
        mlflow.log_metric(
            "time/epoch", (time.time() - start) / 60 / training_args.num_train_epochs
        )  
```

4. After running my bash script, I see some amount of GPU being used (10G/80G) on all of the 6 GPU's, but it hangs after logging this --

```
***** Running training *****
  Num examples = 364978
  Num Epochs = 3
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 48
  Gradient Accumulation steps = 1
  Total optimization steps = 22812
  Number of trainable parameters = 1030852826

  0%|          | 0/22812
```

### Expected behavior

Expected Behaviour is for the training to start and the processes to not hang.
Ok the issue is not related to FSDP per say. Looks like something related to communication between the GPU's. I fixed it by 
 modifying an environment variable as follows : `export NCCL_P2P_DISABLE=1` 
Closing!