Fix XGLM loss computation (PyTorch and TensorFlow)

[damianoamatruda]

What does this PR do?

This PR fixes the loss computation for XGLM in both PyTorch and TensorFlow implementations.

The labels were shifted by one and the padding token was appended, causing artificial loss contributions, inconsistencies between non-padded and right-padded sequences, and potential bias toward predicting padding tokens.

The updated implementations ignore the last logit and do not append the padding token to the labels, aligning with the behavior in GPT-2 and other models.

The logic of the computation was first identified in #22540, where it was ported from the PyTorch implementation to the TensorFlow one for consistency. In this PR I've reverted the TensorFlow implementation to its previous valid behavior and I've updated the PyTorch implementation to match it.

I've also added XGLM tests to ensure that the losses of non-padded and padded inputs match.

This bug was discovered in a joint project while collaborating with @mdrpanwar and @ayushkumartarun.

Who can review?

@Rocketknight1 @gante @ArthurZucker

[Rocketknight1]

Hi @damianoamatruda, yes, the original code is incorrect! However, a simpler fix would be to change the label padding value to -100, which indicates masked positions for our loss computations. Padding/shifting labels is preferable to shifting logits, because logits are a much larger tensor and can carry gradient, so speed + memory usage are affected if we perform lots of operations on them, especially during training.

[damianoamatruda]

Hi @Rocketknight1, thank you for the clear explanation!

I've updated the PR to shift only the labels, as previously done, and replaced the padding token with the mask value -100.

I've also updated the PyTorch test to match the changes introduced in the newly merged PR #35659.

[Rocketknight1]

Yes, LGTM now! cc @ArthurZucker for core maintainer review

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Rocketknight1]

run-slow: xglm

[github-actions]

This comment contains run-slow, running the specified jobs: ['models/xglm'] ...

[Rocketknight1]

Hi @damianoamatruda I'm seeing some failures in the slow tests for XGLM, can you take a look? You can check the CI logs, or to run slow tests locally, you can do something like RUN_SLOW=1 pytest tests/models/xglm/test_modeling_xglm.py

[damianoamatruda]

Hi @Rocketknight1, I took a look at the errors, which were related to XGLM and similar models but weren't connected to the loss computation. I fixed them by taking inspiration from models/t5/modeling_tf_t5.py. Could you give me your opinion?

Now, however, with the latest rebase, there are failed tests that aren't related to XGLM. Can you do something about it?

[Rocketknight1]

Hi @damianoamatruda, I'm not sure exactly what's causing that! It's likely those tests were just flaky on a past commit - can you try rebasing again? If they still won't go away then I'll see if we can actually fix or skip them on main.

[damianoamatruda]

@Rocketknight1, I rebased and the test tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py::Qwen2_5_VLModelTest::test_generate_from_inputs_embeds_1_beam_search failed again.

[Rocketknight1]

Yeah, that test is a problem on main right now. Please leave it a couple of days and then check back in - hopefully we'll have resolved it by now!

[Rocketknight1]

Tests are finally green! Pinging @Cyrilvallez for core maintainer review

[damianoamatruda]

Great!

[Cyrilvallez]

Hey! All LGTM concerning the loss part!

However, I must say that I am skeptical concerning the change in set/get embeddings. It looks like we are changing input type to the functions here (layer vs underlying layer data), which may be breaking for current code. Moreover, the failing test explicitly states that it is expected to fail (and it was never fixed).
TLDR I'd rather we revert the part on embeddings, and keep the loss part 🤗

[damianoamatruda]

Hi @Cyrilvallez, done!

The tests now pass without requiring the commits for the embeddings. How did you fix/disable the failing tests?

Thank you for the review.

[Cyrilvallez]

Thanks for reverting! The failing tests were slow test triggered by the github-actions, they are not run by the usual CI which is why you cannot see them now!

[damianoamatruda]

Is there anything else to do or is everything okay?

[ArthurZucker]

Let's go! @Cyrilvallez has a conference this week hahah sorry 🤗

[ArthurZucker]

BTW you need to resolve confilcts (probably no changes needed on the modeling non tf side no?

[Cyrilvallez]

Hey @damianoamatruda! Indeed very sorry, I had a lot going on this week! As you can see, in the meantime xglm got the loss refactor incorporated, which automatically fixed the issue at hand in Pytorch. The change to the pytorch modeling should not be needed anymore. Very happy to add the test and the change to tensorflow file though!

[damianoamatruda]

@ArthurZucker, @Cyrilvallez, no problem, I understand your commitments 🤗

[damianoamatruda]

Refactor #35875 moved the loss computation for XGLM into a dedicated function xglm_cross_entropy_loss, but the invalid logic remains: it still appends the padding token to the labels, causing this PR's test XGLMModelLanguageGenerationTest::test_loss_with_padding to fail. Removing the dedicated function fixes the issue.

[Cyrilvallez]

Oh indeed, did not notice that #35875 created a dedicated function to ensure BC, but BC was wrong!

All right, LGTM! Thanks a lot for the fix!! 🤗

[damianoamatruda]

Thank you all, it's been a pleasure for me! 🤗