Skip to content

[TF] Fix Tensorflow XLA Generation on limited seq_len models#33903

Merged
ArthurZucker merged 3 commits intohuggingface:mainfrom
vasqu:fix-xla-generation-on-fixed-seq-lens
Oct 5, 2024
Merged

[TF] Fix Tensorflow XLA Generation on limited seq_len models#33903
ArthurZucker merged 3 commits intohuggingface:mainfrom
vasqu:fix-xla-generation-on-fixed-seq-lens

Conversation

@vasqu
Copy link
Contributor

@vasqu vasqu commented Oct 3, 2024

What does this PR do?

Fixes xla generation with models with a limited seq_len that they can generate. It has been recently (re)discovered in #33298, which results in slow runs to fail as we have indexing issues based on absolute positional embeddings (i.e., an index out-of-bounds error). The issue can only be observed on CPUs and not on GPUs (idk why but the error is just silently swallowed).

The solution is to extend the embedding not only by max_new_tokens but also by the current seq_len at hand. Atm, seq_len hasn't been considered yet.

No fixes but a ref to #33298
(possibly adding an issue if wanted)

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?

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.

@Rocketknight1 @amyeroberts

inputs = inputs_dict["input_ids"] if is_input_ids else inputs_dict["input_features"]

# fix config for models with additional sequence-length limiting settings
seq_len = inputs.get_shape()[1]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if input_features also has its sequence length at 1. Just submitting this quickly before I don't have time.

@Rocketknight1
Copy link
Member

@vasqu I suspect the reason that these errors are silent on GPU is that tf.gather() does not do bounds checking on GPU:

Caution: On CPU, if an out of bound index is found, an error is raised. On GPU, if an out of bound index is found, a 0 is stored in the corresponding output value.

tf.gather() is the underlying mechanism used for a lot of things like embedding layers and other index lookups, so this affects a lot of layers.

Either way, reviewing the rest of this now!

@Rocketknight1
Copy link
Member

@vasqu can you push another empty commit with the [run-slow] opt message to trigger the slow CI there?

@vasqu
Copy link
Contributor Author

vasqu commented Oct 3, 2024

Oh wow, thanks for the insights into this. This is rather counter-intuitive to me tho :D Even some archaic cuda / gpu errors would suffice for me but it is what it is.

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

@vasqu
Copy link
Contributor Author

vasqu commented Oct 4, 2024

tests pass :) gentle ping @Rocketknight1

Copy link
Member

@Rocketknight1 Rocketknight1 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

For core maintainers: This generation test fails on some models because it generates past the limit of the model's max_position_embeddings. Because the model is initialized from a config, not a pretrained checkpoint, this PR simply increases config.max_position_embeddings or config.max_target_positions, which seems like the right fix to me.

cc @LysandreJik for core maintainer review!

@vasqu vasqu mentioned this pull request Oct 4, 2024
5 tasks
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks 🤗

@ArthurZucker ArthurZucker merged commit 5ef432e into huggingface:main Oct 5, 2024
@vasqu vasqu deleted the fix-xla-generation-on-fixed-seq-lens branch October 5, 2024 14:43
BernardZach pushed a commit to BernardZach/transformers that referenced this pull request Dec 5, 2024
…ngface#33903)

* fix tf xla generation on limited seq_len models

* [run-slow] opt

* [run-slow] opt
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants