Fix for hardcoded `final_labels` to enable loss calculation in PaliGemma

[ariG23498]

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.

transformers/src/transformers/models/paligemma/modeling_paligemma.py

Line 349 in 87a3518

final_labels = None

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

[amyeroberts]

cc @molbap

[molbap]

Hi @ariG23498 , thanks for the issue! Labels should be formed correctly by the processor if you pass suffix.

processor(images=img, text=["What is that?"], suffix=["a cow!"], return_tensors="pt", padding="longest")

In your case suffix would be example['blip_caption'].

[ariG23498]

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?

[molbap]

Checking, your snippet is

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?

[ariG23498]

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.

[molbap]

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

[ariG23498]

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?

transformers/src/transformers/models/paligemma/modeling_paligemma.py

Line 478 in 6657fb5

logits = logits.float()

[molbap]

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

[ariG23498]

That makes sense. Thanks for the reply!