Fix setting vocab_size when resizing embeddings and using deepspeed zero 3
Author: ppetrushkov
Pull request: yes

# What does this PR do?

Currently there seems to be an issue when training with deepspeed zero stage 3 and resizing the token embeddings. Here is an example I used:

```python
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    HfArgumentParser,
    TrainingArguments,
    DataCollatorForSeq2Seq,
    Trainer,
)

def main():
    parser = HfArgumentParser((TrainingArguments))
    training_args, = parser.parse_args_into_dataclasses()

    model_path = 'facebook/opt-125m'

    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        config = config,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        model_max_length=1024,
    )

    add_new_tokens = True
    if add_new_tokens:
        tokenizer.add_special_tokens({"pad_token": "<pad>",})
        model.resize_token_embeddings(len(tokenizer))
    else:
        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()
```

Executed on at least 2 GPUs with: `deepspeed minimal.py --output_dir output_dir --deepspeed deepspeed_config.json`

where deepspeed_config.json is:

```json
{
  "bf16": {
    "enabled": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": "auto",
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto",
      "warmup_type": "linear",
      "total_num_steps": "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": 2000,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}
```



Currently the `vocab_size` and `config.vocab_size` are set to 0, when using deepspeed zero in this line https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L1537-L1538 . This fails later during training with `RuntimeError: shape '[-1, 0]' is invalid for input of size 804240` when trying to flatten the logits using `vocab_size` shape. Instead it should use approach similar to https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L1555-L1559 to set the correct size.

## Who can review?

@ArthurZucker @pacman100 




Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26387). All of your documentation changes will be reflected on that endpoint.

Hello @ppetrushkov , this merged PR https://github.com/huggingface/transformers/pull/26259 should have fixed this issue. Could you retry from the main branch?

@pacman100 can confirm it works from the main branch. I'll close this PR then.