TF: XLA stable softmax

[gante]

What does this PR do?

As discussed in the thread about XLA problems (#16838), this PR adds a stable wrapper for the softmax operation, and replaces tf.nn.softmax by the wrapped function.

This PR:

• Adds the wrapped softmax, named stable_softmax, in tf_utils.py. Its docstring includes why it is needed and why the new operation is valid;
• Adds tests to the wrapped softmax, including XLA tests;
• Replaces tf.nn.softmax by stable_softmax everywhere except in the doctests (I think it overcomplicates the examples, and no XLA should be needed there);
• Removes the skipIf for XLA tests, as they can now be successfully executed in a CPU.

Closes #16838

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Great find for the bug and thanks a lot for fixing all models!

[ydshieh]

Could we have a test for bacth_size > 1 ?

[gante]

Added batch size > 1 👍

[ydshieh]

Good for me, just left 2 nits.
(didn't check the changes in TFGPT2, TFT5 tests though. Let me know if you prefer me to check those too.)

Thank you, @gante 💯

[Rocketknight1]

This looks good to me! Do you think it would be better to change stable_softmax to only add the offset if we're running on CPU? It makes very little difference either way, but we could hide the complexity of that inside stable_softmax and keep our code paths entirely unchanged on GPU. I'm not certain, though - since it's such a small change maybe we can just do it everywhere.

[ydshieh]

This looks good to me! Do you think it would be better to change stable_softmax to only add the offset if we're running on CPU? It makes very little difference either way, but we could hide the complexity of that inside stable_softmax and keep our code paths entirely unchanged on GPU. I'm not certain, though - since it's such a small change maybe we can just do it everywhere.

Good point! Hope this won't affect tests on GPU (at least not for PT/TF equivalence which use 1e-5). Let's see!

[gante]

@Rocketknight1 @ydshieh if you run the test and print the difference between stable_softmax and tf.nn.softmax, the difference is exactly 0.0 -- I don't think we need to worry about that :D

[ydshieh]

@gante With this, do we still have issues regarding sampling in generate(). Sorry, I didn't really follow that issue about sampling, but would like to know a bit more 😄

[gante]

@ydshieh after this fix, the errors related to generate() are gone -- they were caused by the forward pass in the models, which in turn were caused by the issue this PR solves

[ydshieh]

(I might be completely wrong below)

I could imagine that we (will) have tests like:

• testing non-XLA and XLA generte() that use sampling
  • even with this PR, the differences of output logits between these two might still be as large as, say, 1e-3?
  • if so, the sampling might give different sampling results ..?
  • if not, what's the magnitude of the diff we get after this PR?
• testing PT and TF generte() that use sampling
  • so same potential issue as above ..?

Thanks 🙏

[ydshieh]

OK, I saw your previous comment

I've spun up an Nvidia T4 ( = no tf32 format) and got an error < 1e-5 for all cases

[Rocketknight1]

Based on the testing results, I'm happy for this to be merged now! If this is an XLA bug, though, we should make sure to revert our changes once none of the TF versions we support are affected by it anymore.

Should we add a TODO to the masked_softmax function or a reminder somewhere to make sure that we document why this change is here, and when it can be removed?

[patrickvonplaten]

Great - the solution is clean!

[gante]

@Rocketknight1 added a TODO with instructions related to when to deprecate 👍