Mixtral gets stuck at Loading checkpoint shards.
Author: AdamLouly
Pull request: no
Labels: name: bug

name: Big Model Inference

### System Info

Nightly transformers.
nightly torch

8 GPUs

### Who can help?

When trying to run mixtral using the example in transformers.
it gets stuck at Loading checkpoint shards
at this point:
Loading checkpoint shards:  42%|██████████████████████████████████████████████▋   
                                                             
Noticed this only happens when running on multiple GPUS
### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

run the fine tunning example on transformers

### Expected behavior

stuck at 
Loading checkpoint shards:  42%|██████████████████████████████████████████████▋                                                                

Comments:
Hi @AdamLouly, thanks for raising an issue! 

Could you share a reproducible snippet to show how you're loading the model? 
For the transformers version - could you share the commit you're running on?

cc @ArthurZucker @younesbelkada 

hi @AdamLouly 
As stated by @amyeroberts , a reproducible snippet would be much appreciated! Also how do you run your script? `python xx.py` or `accelerate launch xxx.py` or `python -m torch.distributed.run xx.py` etc. 🙏 

Hi @amyeroberts 
this is how I install. transformers:
pip install git+https://github.com/huggingface/transformers
so basically, latest version from source.

Hey @younesbelkada 
this is the snippet:


`#!/bin/bash

ds_config=$(mktemp --suffix ".json")
echo "The DeepSpeed config is put at $ds_config"

cat <<EOF > $ds_config
{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 20000000,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 20000000,
    "contiguous_gradients": false,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  "zero_allow_untested_optimizer": true,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": "auto",
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },
  "steps_per_print": 2000,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "wall_clock_breakdown": false
}
EOF

num_gpus=8

torchrun --nproc_per_node $num_gpus \
  run_glue.py \
    --model_name_or_path mistralai/Mixtral-8x7B-v0.1 \
    --task_name MRPC \
    --per_device_train_batch_size 1 \
    --do_train \
    --num_train_epochs 5 \
    --output_dir "./results" --overwrite_output_dir \
    --save_strategy "no" \
    --fp16 --max_steps 100 \
    --gradient_accumulation_steps 1 \
    --learning_rate 0.00001 \
    --adam_beta1 0.9 \
    --adam_beta2 0.999 \
    --adam_epsilon 1e-8 \
    --deepspeed $ds_config`

Then I run the script this way : bash run_mixt.sh

Hi @AdamLouly 
Not sure what is wrong here, can you test on one of the latest pypi release to check if this is not a regression? for example you can try: `pip install -U transformers==4.36.2`. Can you also alternatively try with `accelerate`? First run `accelerate config`, select DeepSpeed and run `accelerate launch run_glue.py`

Hey @younesbelkada I tested both and still having the same issue.

Have you managed to reproduce it using the example above?


This might be related to NCCL communications no? We have a lot of issues likes this `NCCL_P2P_DISABLE=1` should fix it. 

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.

try this:
om transformers import AutoTokenizer, AutoModelForCausalLM
repo_id = "Open-Orca/Mistral-7B-OpenOrca"
model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype="auto", low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained(repo_id, torch_dtype="auto", low_cpu_mem_usage=True)