Error when changing vocab size when fine tuning llama-vision
Author: Ssukriti
Pull request: no
Labels: name: bug

### System Info

**- `transformers` version: 4.49.0**
- Platform: Linux-5.14.0-284.73.1.el9_2.x86_64-x86_64-with-glibc2.34
- Python version: 3.12.5
- Huggingface_hub version: 0.28.1
- Safetensors version: 0.5.2
- Accelerate version: 1.0.1
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.4.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 distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

@amyeroberts, @qubvel 

### Information

- [x] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

while fine tuning model : [Llama-3.2-11B-Vision-Instruct](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct)

if I add new tokens and resize embeddings, I get an error while fine tuning as the modified vocab size does not make it to the language model class. **I have done an analysis of the error and will provide more details more, and a potential fix.** Open to feedback . 

Script to reproduce for illustration: 
```

from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
from transformers import AutoProcessor, AutoModelForVision2Seq
import torch
from typing import Dict
import transformers
import math 

datapath = "HuggingFaceH4/llava-instruct-mix-vsft"
dataset = load_dataset(datapath, split="train[:100]")

model_path = "Llama-3.2-11B-Vision-Instruct/"
processor = AutoProcessor.from_pretrained(model_path)
model = AutoModelForVision2Seq.from_pretrained(model_path, torch_dtype=torch.bfloat16)
tokenizer = processor.tokenizer
    

# code to resize embeddings that calls model.resize_token_embeddings
def tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
    multiple_of: int = 1,
) -> dict:
    """Resize tokenizer and embedding.
    Args:
        special_tokens_dict: Dict containing special tokens to be added.
        tokenizer: transformers.PreTrainedTokenizer.
        model: transformers.PreTrainedModel
        multiple_of: int , embeddings are resized to multiple of this.
    Return:
        dict: Metadata on number of added tokens
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    embedding_size = int(multiple_of * math.ceil(len(tokenizer) / multiple_of))
    num_new_tokens = num_new_tokens + embedding_size - len(tokenizer)
    model.resize_token_embeddings(embedding_size)
    print("\n resized vocab size ", model.config.get_text_config().vocab_size)


# collate function
def collate_fn(examples):
    texts = [processor.apply_chat_template(example["messages"], tokenize=False) for example in examples]
    images = [example["images"] for example in examples]

    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

    labels = batch["input_ids"].clone()
    if processor.tokenizer.pad_token_id is not None:
        labels[labels == processor.tokenizer.pad_token_id] = -100
    # Ignore the image token index in the loss computation (model specific)
    image_token_id = processor.tokenizer.convert_tokens_to_ids(
        processor.image_token
    )
    labels[labels == image_token_id] = -100
    batch["labels"] = labels
    return batch

peft_config = LoraConfig(
        r=8,
        lora_alpha=8,
        lora_dropout=0.1,
        target_modules=["q_proj", "v_proj"],
    )

# setting a new UNK token (what token doesnt matter here, just an example)
special_tokens_dict = {}
if tokenizer.unk_token is None:
    print("UNK token set to default, missing in tokenizer")
    special_tokens_dict["unk_token"] = "<UNK>"


added_tokens_dict = tokenizer_and_embedding_resize(
        special_tokens_dict=special_tokens_dict,
        tokenizer=tokenizer,
        model=model,
        multiple_of=1,
)


# Configure training arguments using SFTConfig
training_args = SFTConfig(
    output_dir="./checkpoints/llama_llava",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=1e-4,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,
    push_to_hub=False,
    report_to="none",
    remove_unused_columns=False,
    gradient_checkpointing=True,
    dataset_text_field="messages",
    dataset_kwargs={"skip_prepare_dataset": True},
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=collate_fn,
    peft_config=peft_config,
    processing_class=processor.tokenizer,
)


trainer.train()

# RUN WITH
# accelerate launch --num_processes=2 --config_file /app/accelerate_fsdp_defaults.yaml llama_llava_dataset.py
``` 

Error: 
```
[rank1]:   File "/usr/local/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py", line 1967, in forward
[rank1]:     loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/site-packages/transformers/loss/loss_utils.py", line 43, in ForCausalLMLoss
[rank1]:     logits = logits.view(-1, vocab_size)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: shape '[-1, 128256]' is invalid for input of size 445311776
```

**Detailed explanation after debugging:**

1. 128256 is the original vocab size . After resizing the model.config.get_text_config().vocab_size is 128258
2. the modified vocab size is applied to `MllamaForConditionalGeneration` class. The embeddings of language model are resized as expected. However the vocab size in language model class being pass to loss function is not modified. Thus the error comes from https://github.com/huggingface/transformers/blob/main/src/transformers/models/mllama/modeling_mllama.py#L1966 where self.vocab size is still  128256
3. the shape of logits are modified as output embeddings are resized, but vocab size is not , causing a mismatch
**Hence, the language model's config needs to be modified with updated vocab size**

### Expected behavior

The language model's vocab size and embeddings should be resized / updated. The embeddings are resized as expected, but vocab size is not.

Tuning should not result in any error when adding tokens

Comments:
potential fix https://github.com/huggingface/transformers/pull/36591 . This fixes the issue

There are several ways to do this though and open to discussion. PR in draft for discussion.

@zucchini-nlp @SunMarc Thanks!

Thanks for reporting! I agree that the loss function should be getting config's vocab size, which is the one modified when we resize embeddings. Left a comment under PR :)

We might also want to check other models if they use the correct vocab size, after agreeing on the fix

closing as completed