Assertion error when using Trainer & Deepspeed stage 3 with `model.resize_token_embeddings`
Author: kai01ai
Pull request: no

### System Info

- `transformers` version: 4.33.0
- Platform: Linux-3.10.0-1160.83.1.el7.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.9
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.2
- Accelerate version: 0.22.0
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: DEEPSPEED
	- mixed_precision: bf16
	- use_cpu: False
	- debug: False
	- num_processes: 64
	- machine_rank: 0
	- num_machines: 8
	- main_process_ip: 127.0.0.1
	- main_process_port: 29500
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- deepspeed_config: {'deepspeed_multinode_launcher': 'standard', 'gradient_accumulation_steps': 1, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': True, 'zero3_save_16bit_model': True, 'zero_stage': 3}
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- 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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### 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, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I've encountered an issue while training using the combination of Trainer & Deepspeed stage 3. When invoking model.resize_token_embeddings, an AssertError arises during training. This was not an issue in transformers version 4.31.0. However, for versions > 4.31.0 and in the main branch, this problem persists. I suspect this might be related to PR https://github.com/huggingface/transformers/pull/25394

```
  File "/root/miniconda3/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 310, in fetch_sub_module
    return func(*args, **kwargs)
  File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 310, in fetch_sub_module
    assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
    AssertionErrorassert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 292, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0, 'shape': (0,), 'ds_shape': (0, 4096), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {453}, 'ds_tensor.shape': torch.Size([0])}
```
code:
```python test.py
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    HfArgumentParser,
    TrainingArguments,
    DataCollatorForSeq2Seq,
    Trainer,
)

def main():
    parser = HfArgumentParser((TrainingArguments))
    training_args, = parser.parse_args_into_dataclasses()

    model_path = '/path/to/Llama-2-7b-hf'

    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        config = config,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_path, 
        use_fast=False,
        model_max_length=1024,
    )
    
    add_new_tokens = True
    if add_new_tokens:
        # deepspeed AssertionError
        tokenizer.add_special_tokens({"pad_token": "<pad>",})
        model.resize_token_embeddings(len(tokenizer))
    else:
        # it works
        tokenizer.pad_token = tokenizer.eos_token

    from datasets import Dataset
    def gen():
        for _ in range(100):
            yield {"input_ids": [1, 2, 3], "labels": [1, 1, 1]}
    datasets = Dataset.from_generator(gen)
    datasets.set_format('pt')

    trainer = Trainer(
        model=model,
        args=training_args,
        tokenizer=tokenizer,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, max_length=tokenizer.model_max_length),
        train_dataset=datasets,
    )
    trainer.train()


if __name__ == "__main__":
    main()
```
scripts:
```shell
deepspeed test.py \
    --deepspeed configs/zero3_hf.conf \
    --output_dir output/test/ \
    --do_train \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --evaluation_strategy "no" \
    --report_to "none" \
```
deepspeed config
```conf
{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 1e5,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```

### Expected behavior

no error

Comments:
Meet the same error. Exactly the same situation: transformers **> 4.31.0** + deepspeed zero3 v0.10.2 + `resize_token_embeddings`.

same issue for `transformers >= 4.32.0`

Try to fix this by https://github.com/huggingface/transformers/pull/26024