`from_pretrained` is extremely slow when deepspeed zero3 is enabled
Author: Jingru
Pull request: no

### System Info

pytorch: 2.0.1+cu118
transformers: 4.33.3
deepspeed: 0.12.5

### Who can help?

@ArthurZucker  @younesbelkada @pac

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

1. Run command `torchrun --nnodes 1 --nproc-per-node 8 --rdzv-endpoint=localhost:35000 test.py`

And my script `test.py` as follows:

```
import deepspeed

from transformers.deepspeed import HfDeepSpeedConfig
from transformers import AutoModelForCausalLM


deepspeed.init_distributed()

ds_config = {
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 4,
    "steps_per_print": 10,
    "zero_optimization": {
        "stage": 3,
        "offload_param": {"device": "cpu"},
        "offload_optimizer": {"device": "cpu"},
        "stage3_param_persistence_threshold": 10000.0,
        "stage3_max_live_parameters": 30000000.0,
        "stage3_prefetch_bucket_size": 30000000.0,
        "memory_efficient_linear": False,
    },
    "fp16": {"enabled": True, "loss_scale_window": 100},
    "gradient_clipping": 1.0,
    "prescale_gradients": False,
    "wall_clock_breakdown": False,
    "hybrid_engine": {
        "enabled": True,
        "max_out_tokens": 512,
        "inference_tp_size": 1,
        "release_inference_cache": False,
        "pin_parameters": True,
        "tp_gather_partition_size": 8,
    },
}

dschf = HfDeepSpeedConfig(ds_config)

model = AutoModelForCausalLM.from_pretrained(
    "../llama_actor", from_tf=False, trust_remote_code=False
)

```
In addition, the pretrained model is saved by `transformers==4.31.0`.

2. This command hangs for over 1800s, and encountered nccl timeout error.


### Expected behavior

Model is loaded in a few minutes and this command should not hang.

Comments:
cc @pacman100 

I did some research, and found that all the ranks except rank0 tried to load model weight tensors to `meta device`.  It looks like this behavior is extremely slow. 
https://github.com/huggingface/transformers/blob/bffac926ca6bc6c965a92bfbfd00c567a2c0fb90/src/transformers/modeling_utils.py#L485

Hello, the state dict is loaded loaded on rank 0 to avoid excessive memory usage and loading the model on meta on other devices should be the fastest as no materialization takes places. 

Code changes to your example as i can't access the private model:
```diff
import deepspeed

from transformers.deepspeed import HfDeepSpeedConfig
from transformers import AutoModelForCausalLM


deepspeed.init_distributed()

ds_config = {
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 4,
    "steps_per_print": 10,
    "zero_optimization": {
        "stage": 3,
        "offload_param": {"device": "cpu"},
        "offload_optimizer": {"device": "cpu"},
        "stage3_param_persistence_threshold": 10000.0,
        "stage3_max_live_parameters": 30000000.0,
        "stage3_prefetch_bucket_size": 30000000.0,
        "memory_efficient_linear": False,
    },
    "fp16": {"enabled": True, "loss_scale_window": 100},
    "gradient_clipping": 1.0,
    "prescale_gradients": False,
    "wall_clock_breakdown": False,
    "hybrid_engine": {
        "enabled": True,
        "max_out_tokens": 512,
        "inference_tp_size": 1,
        "release_inference_cache": False,
        "pin_parameters": True,
        "tp_gather_partition_size": 8,
    },
}

dschf = HfDeepSpeedConfig(ds_config)

- model = AutoModelForCausalLM.from_pretrained(
-     "../llama_actor", from_tf=False, trust_remote_code=False
- )
+ model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
```

launch command on 2 GPUs:  
```
time torchrun --nnodes 1 --nproc-per-node 2 --rdzv-endpoint=localhost:35000 issue_28139.py
```

Result:
```
[2023-12-19 15:08:18,112] torch.distributed.run: [WARNING] 
[2023-12-19 15:08:18,112] torch.distributed.run: [WARNING] *****************************************
[2023-12-19 15:08:18,112] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2023-12-19 15:08:18,112] torch.distributed.run: [WARNING] *****************************************
[2023-12-19 15:08:21,074] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-12-19 15:08:21,074] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/transformers/utils/hub.py:123: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.
  warnings.warn(
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/transformers/utils/hub.py:123: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.
  warnings.warn(
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations
  warnings.warn(
[2023-12-19 15:08:23,368] [INFO] [comm.py:637:init_distributed] cdb=None
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations
  warnings.warn(
[2023-12-19 15:08:23,499] [INFO] [comm.py:637:init_distributed] cdb=None
[2023-12-19 15:08:23,499] [INFO] [comm.py:668:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2023-12-19 15:08:31,060] [INFO] [partition_parameters.py:348:__exit__] finished initializing model - num_params = 291, num_elems = 6.74B
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.86s/it]
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.88s/it]

real	0m21,615s
user	0m19,894s
sys	0m11,505s
```

So, checkpoint loaded and sharded across the 2 GPUs in 21 seconds for a 7B model which is inline with the expected behaviour.



I tried model "openlm-research/open_llama_7b_v2" and this model could be loaded as expected.

Does this mean that there is something abnormal with my finetuned model?

I tried to override `map_location` mentioned above to 'cpu' for all the ranks, and this script could load my model in 300s. But the results of `model.generate` are nonsense.

Besides, I tried to downgrade transformers to 4.31.0 and deepspeed to 0.11.1, and everything is fine. Model can be loaded in a few minutes and generated texts are normal.

hi @Jingru, same problem, did u solve this?

> hi @Jingru, same proble, did u solve this?

I load the old checkpoint without deepspeed and `save_pretrained` again. The new checkpoint can be loaded with zero3 normally.

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.