Bug in `PreTrainedModel.resize_token_embeddings` When Using DeepSpeed Zero Stage 3
Author: sinamoeini
Pull request: no

### System Info

transformers version: 4.31.0
Platform: Linux 5.4.238-148.346.amzn2.x86_64
Python version: 3.8.10
Huggingface_hub version: 0.14.1
Safetensors version: 0.3.1
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

### Information

- [ ] The official example scripts
- [X] 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

This is a simple test to highlight this inconsistency. Here is brief description of what test script does:

* Starts deepspeed 
* Loads a pretrained model
* Using gather gets the weights of first 50 embeddings on each device and stores them in a local tensor
* Reduce the number of embeddings to 50 by using `PreTrainedModel.resize_token_embeddings`
* gets the embedding weights again (note that at this point they are not ds pararmeters anymore)
* Checks the result on each device to see if it matches what we recorded earlier

The script is executed on a multi gpu node as follows
```
deepspeed test.py 
```
Where the contents of `test.py` are

```
from transformers import (
    TrainingArguments,
    AutoModelForCausalLM,
    set_seed,
)
import os
import deepspeed

def main() -> None:
    set_seed(0)
    # enable deepspeed stage 3
    training_args = TrainingArguments(output_dir="dummy", remove_unused_columns=False, deepspeed="zero3.json")
    # load pretrained model
    model_path = "openlm-research/open_llama_3b"
    model = AutoModelForCausalLM.from_pretrained(model_path)

    # store first 50 embeddings locally in ref
    with deepspeed.zero.GatheredParameters(list(model.lm_head.parameters())):
        ref = model.lm_head.weight.data[:50, :].clone()

    # reduce embeddings to 5, using resize_token_embeddings
    model.resize_token_embeddings(50)

    # check if the embeddings match what we recorded earlier on each device
    # note that after resizng, resize_token_embeddings does not convert the embedding layers to ds parameters 
    rank = int(os.environ["RANK"])
    sanity = all((ref == model.lm_head.weight.data).reshape(-1).tolist())
    print(f"{rank}: sanity pass: {sanity}")


if __name__ == "__main__":
    main()
```

And contents of `zero3.json` are

```
{
  "train_micro_batch_size_per_gpu": "auto",
  "train_batch_size": "auto",
  "zero_allow_untested_optimizer": true,
  "gradient_clipping": "auto",
  "gradient_accumulation_steps": "auto",
  "bfloat16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": false,
    "overlap_comm": true,
    "allgather_bucket_size": 1e8,
    "reduce_bucket_size": 2e8,
    "stage3_max_live_parameters": 0.7e8,
    "stage3_param_persistence_threshold": 5e6,
    "stage3_gather_fp16_weights_on_model_save": true
  },
  "activation_checkpointing": {
    "partition_activations": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": 100,
    "cpu_checkpointing": false
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "weight_decay": "auto",
      "betas": [
        0.9,
        0.999
      ],
      "eps": "auto",
      "lr": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto",
      "total_num_steps": "auto"
    }
  }
}
```




### Expected behavior

It is expected that each device would return True. But as of now only the 0th device has the correct value. 

Note that in majority of the cases `deepspeed.initialize` is called by Trainer right after `resize_token_embeddings`, where rank 0 values would be scattered and we will have consistency. However if an operation happens in between there will be no consistency


Comments:
Hi! Would it possible for you to do `resize_token_embeddings` without DeepSpeed, save the model, and load the new model in the script where you use DeepSpeed.

This might be easier and quicker in terms of solution/workaround (if it works).

Hi, thanks for the suggestion. I have RCed this and have a nonhacky solution that works nicely. I will create a PR in the next two days to resolve this.

Thank you @sinamoeini for fixing this issue 🤗

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.