XLA train step fixes
Author: Rocketknight1
Pull request: yes

This PR makes a bunch of changes to the TF codebase to improve XLA support, in preparation for our upcoming big TF release. The goal is to allow users to use `jit_compile` on the vast majority of our models, which should yield large performance improvements for TF. In particular:

- Rewrites to the `train_step` and `test_step` so that any mutable Python input dicts are not modified in the step. This was a bad idea anyway, but it causes particular problems with XLA, which is very functional and hates side effects, like JAX.
- Rewrites to the common `hf_compute_loss` functions to ensure that static shapes are maintained throughout, so that XLA compilation is possible.
- Add a test to ensure that we can still fit models when XLA compilation is used. XLA compilation is quite expensive, which makes this test quite slow, so it's restricted to `core` models for now and tagged as `@slow`.

Left to do:
- [x] Fix XLA-incompatible model-specific  `hf_compute_loss` functions. On a quick search it looked like there were 4-5 of these, so it shouldn't take too long. Any use of `tf.boolean_mask` is a surefire sign that XLA compilation will break, because output shapes become data-dependent.
- [x] See if there's a way to test non-core models for XLA fit support without crippling performance. (No, but we're using the XLA losses in non-XLA tests by default, so that partially tests it for all models)

Comments:
_The documentation is not available anymore as the PR was closed or merged._

I'd be interested in having @ydshieh's review as well

> If I understand correctly, this changes the loss returned by TensorFlow models from a matrix to a vector, which is obviously breaking. While I don't know how many TensorFlow users rely on the current structure of the loss, we at least need to have a flag (probably `use_xla=False`) to enable the previous behavior for users who relied on it.
> 
> Could you confirm first that my understanding is correct?

I believe both prev. and current version return a vector. The difference is on the size: 

- prev: number of active tokens (non-padding tokens)
- now: batch size

@ydshieh I was completely wrong earlier - `SparseCategoricalCrossentropy` only returns `nan` for invalid labels when running on GPU! On CPU, inputs are validated and TensorFlow throws an error.

I'll rewrite my loss functions to not depend on that behaviour, and change the loss computation tests to mask some positions to ensure that gets tested, so I don't miss anything like this in future.

@Rocketknight1 

https://github.com/huggingface/transformers/blob/f17136c80dfec2a78890d012105634079531dcd9/src/transformers/modeling_tf_utils.py#L211-L213

As in an earlier comment, I think this loss value is incorrect. Imagine we have 2 sequences of length 100.
- 1st sentence: 1 active token + 99 pad tokens (somehow non-sense 😄 )
- 2nd sentence: 20 active token + 80 pad tokens

In this latest version, the unique token in sentence 1 get an weight (when computing the loss) 20 times larger than each token in the 2nd sentence. (As you first average the loss along sequence dimension).

Furthermore, this doesn't correspond to PyTorch's computation, which leads to test failures (I didn't check in detail if this is the cause, but I believe it is).

**Q: Is there any reason we don't want to sum each token's loss value?**

cc @gante  @patrickvonplaten @sgugger 



Hi @ydshieh I'm sorry, I think you're right there! Let me investigate and see if I can make a PR to weight tokens properly, which should hopefully resolve the issue.

@patrickvonplaten Agreed! I fixed that in https://github.com/huggingface/transformers/pull/18013