Fix for hardcoded `final_labels` to enable loss calculation in PaliGemma
Author: ariG23498
Pull request: no

I have encountered an issue with the current implementation where `final_labels` is hardcoded to `None`, which prevents the loss calculation even when `labels` is provided and not `None`.

https://github.com/huggingface/transformers/blob/87a351818e6c680282d86a00b61fe4c55b41477d/src/transformers/models/paligemma/modeling_paligemma.py#L349

Proposed Fix: Assign `labels` to `final_labels`. This way, `final_labels` can be `None` or have some value based on `labels`.

I have created a working fork demonstrating this fix, and it successfully generates loss:

Fork: https://github.com/ariG23498/transformers/tree/aritra/ft-paligemma
Colab Notebook: https://colab.research.google.com/drive/1W27YMyyQG2ynUo6BYvru2xkzCS8ff83m?usp=sharing

Comments:
cc @molbap 

Hi @ariG23498 , thanks for the issue! Labels should be formed correctly by the processor if you pass `suffix`. 
```python
processor(images=img, text=["What is that?"], suffix=["a cow!"], return_tensors="pt", padding="longest")
```
In your case suffix would be `example['blip_caption']`. 

Hey @molbap

Thanks for the tip! This greatly helps with the label creation. But I am afraid this does not mitigate the issue. Here is the colab notebook that shows that the `outputs.loss` is still `None`. 
Colab: https://colab.research.google.com/gist/ariG23498/32abc9caedcec1de5d8ad64e58c77d55/paligemma-ft.ipynb

Am I doing something wrong here?

Checking, your snippet is 
```python
def collate_fn(examples):
    images = [example["image"].convert("RGB") for example in examples]

    text = ["Caption the image." for _ in examples]
    suffix = [example["blip_caption"] for example in examples]

    inputs = processor(
        images=images,
        text=text,
        suffix=suffix,
        return_tensors="pt",
        padding="longest",
    )

    labels = inputs["input_ids"].clone()
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == image_id] = -100
    inputs["labels"] = labels
    inputs = inputs.to(torch.bfloat16).to(device)
    return inputs
```

it looks like your collate_fn is recreating the labels: when passing `suffix` as you do, `inputs['labels']` should already exist and mask out the entire prefix, pad tokens, and image tokens with -100. Maybe you aren't using ignore_index on the prompt?

You were right!

Using your `collate_fn` and sending the entire inputs dictionary now triggers the loss computation.

Thanks for the help!

Please feel free to close this issue.

Thank you for your feedback! Feel free to reopen if you find another issue!

Hey @molbap 

I was wondering whether typecasting `logits` to `float32` has any side effects if I try to load a smaller precision model (bfloat16 or float16) and then fine tune that?

https://github.com/huggingface/transformers/blob/6657fb5fed3f94812ef97dda16e91eee77f03d61/src/transformers/models/paligemma/modeling_paligemma.py#L478



@ariG23498 good question, logits must be cast to float32 since float16 would cause numerical errors. see https://github.com/huggingface/transformers/pull/29285 (same logic with .float() is present in cohere, gemma, llama, mistral etc for this reason)

That makes sense. Thanks for the reply!