Assisted model doesn't seem to be working for Meta-Llama-3-8B
Author: jivanph
Pull request: no

### System Info

- `transformers` version: 4.41.0.dev0
- Platform: Linux-5.10.214-180.855.amzn2int.x86_64-x86_64-with-glibc2.2.5
- Python version: 3.8.6
- Huggingface_hub version: 0.23.0
- Safetensors version: 0.4.1
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.1+cu121 (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?

@ArthurZucker @younes

### 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
```python3
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = 'meta-llama/Meta-Llama-3-8B'
assistant_checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
device = "cpu"

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(**inputs, do_sample=True, assistant_model=assistant_model,num_assistant_tokens=3,num_assistant_tokens_schedule='constant')
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```
The error I get is:

---------------------------------------------------------------------------
IndexError: index out of range in self


### Expected behavior

I would expect a decoded sequence to be printed. But I get an error about an index out of range.

Comments:
```python3

In [3]: assistant_model.config
Out[3]: 
LlamaConfig {
  "_name_or_path": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "hidden_act": "silu",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "intermediate_size": 5632,
  "max_position_embeddings": 2048,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 22,
  "num_key_value_heads": 4,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 10000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.41.0.dev0",
  "use_cache": true,
  "vocab_size": 32000
}

In [4]: model.config
Out[4]: 
LlamaConfig {
  "_name_or_path": "meta-llama/Meta-Llama-3-8B-Instruct",
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.41.0.dev0",
  "use_cache": true,
  "vocab_size": 128256
}
```
the vocab size is very different 😉 you need to resize the embedding or choose a better assistant that was trained on the same vocab size! 
Closing as this is just that the assitant model is not correct! 

Thanks, that makes a lot of sense!