FSDP hangs before training starts
Author: agneet42
Pull request: no

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

Comments:
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!