Delay float32 upcast in ForCausalLMLoss after filtering ignore_index

[starcatmeow]

What does this PR do?

This PR implements the optimization discussed in #38452, originally proposed by @harshit2997.
Thanks for the original suggestion and discussion.

• Move the float32 upcast in ForCausalLMLoss to after filtering out ignore_index labels.
• Ensures only relevant logits are upcasted, reducing VRAM usage without affecting correctness.

Fixes #38452

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. Memory saving by upcasting logits for only non-ignored positions #38452
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

• @ArthurZucker
• @Rocketknight1
• @harshit2997

[Rocketknight1]

Hi @starcatmeow, I just took a look and I'm not sure we can accept this! Although I thought it was a good optimization in the issue at #38452, this is because I misread - I thought the code was already selecting masked labels, it was just doing it after casting to float32 instead of before. In that case, doing the select first would save memory with no issues.

However, adding a mask select step when we didn't use one before introduces some problems - in particular, it makes the sizes of the output tensors data-dependent, which can force recompilations with less efficient dynamic shapes. I'm not sure it's worth it for the memory saving here! cc core maintainers @ArthurZucker @Cyrilvallez for their opinion

[Cyrilvallez]

Interesting idea, I'm not sure how often we expect to see ignore_index in the labels, cc @ArthurZucker do you have an idea? It could actually save a LOT of memory as vocabularies tend to be large now (200k+)

Mathematically, it's fully equivalent

[harshit2997]

Thanks @starcatmeow for picking this change up. @Cyrilvallez my main motivation behind proposing the change was to account for fine-tuning cases where one doesn't want to fine tune on prompt tokens. In addition, won't this also help with getting rid of padding tokens before the upcast which is a common scenario?