Fix use_cache for xla fsdp

[alanwaketan]

What does this PR do?

use_cache cannot be used with gradient checkpointing. In PyTorch/XLA, we have to rely on our own gradient checkpointing function instead of the upstream one. Somehow, transformers regress and couldn't recognize our gradient checkpointing anymore. This PR fixes it.

Fixes #30155

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.
• 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?

@muellerzr @amyeroberts

[alanwaketan]

I don't know how to fix the issue with the use_cache parameter in the modeling code. Is that widely used?

[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.

[amyeroberts]

Hi @alanwaketan, thanks for handling this!

I think this is OK - it's consistent with logic in transformers modeling code, @muellerzr to confirm it's fine for trainer.

I don't know how to fix the issue with the use_cache parameter in the modeling code. Is that widely used?

@alanwaketan Could you specify which modeling code you're referring to?

Re the failing tests - there was a fix pushed on main. Rebasing should resolve

[alanwaketan]

Thanks, @amyeroberts. In most of the modeling code, I saw this use_cache is passed as a parameter:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L949

And then there is this check happened in the modeling code:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L967-L971

Since we currently still cannot directly use the upstream grad ckpt, we cannot relay on the gradient_checkpointing flag and then reuse this logic. And I also don't want to add a new flag and modify all modeling codes. lol

[muellerzr]

Agreed should be fine to me. Thanks!

[alanwaketan]

Thanks, @muellerzr. Can someone help me merge this?

[amyeroberts]

I can merge. Thanks again for fixing this!

[alanwaketan]

Thanks, @amyeroberts