Skip to content

[PyTorch/XLA] Fix extra TPU compilations introduced by recent changes#29158

Merged
amyeroberts merged 4 commits into
huggingface:mainfrom
alanwaketan:alanwaketan/fix_tpu
Mar 13, 2024
Merged

[PyTorch/XLA] Fix extra TPU compilations introduced by recent changes#29158
amyeroberts merged 4 commits into
huggingface:mainfrom
alanwaketan:alanwaketan/fix_tpu

Conversation

@alanwaketan

@alanwaketan alanwaketan commented Feb 21, 2024

Copy link
Copy Markdown
Contributor

What does this PR do?

This PR tries to fix some extra TPU compilations caused by recent HF changes.

  1. PyTorch/XLA doesn't support sdpa yet. So we need to set the default attention implementation to eager.
  2. tensor.item() will trigger tpu graph synchronization. We should avoid using it in the training loop.

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?

@ArthurZucker @younesbelkada

ArthurZucker
ArthurZucker reviewed Feb 21, 2024
View reviewed changes

@ArthurZucker ArthurZucker left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for finding these!

Comment thread src/transformers/trainer.py Outdated

@ArthurZucker ArthurZucker Feb 21, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

great catch !

@muellerzr muellerzr Mar 1, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Can we do this only if we're on a TPU? Otherwise we'll eventually get back to this issue: #29207, which is working on being solved in this PR: #29212

Comment thread src/transformers/modeling_utils.py Outdated

@ArthurZucker ArthurZucker Feb 21, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

This restricts using the model on CPU when is_torch_tpu_available() so a bit too restrictive maybe?

@alanwaketan alanwaketan Feb 21, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Do you have a better idea how to fix this? I'm not familiar with this part of HF....

@ArthurZucker ArthurZucker Feb 23, 2024 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think we should check if the torch being used is xla? something like if is_xla_available?

@alanwaketan alanwaketan Mar 13, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Sorry for late reply. I was busy with something else. I guess this is exactly what my change does?

@muellerzr muellerzr mentioned this pull request Feb 29, 2024
Merged
@alanwaketan

alanwaketan commented Mar 13, 2024

Copy link
Copy Markdown
Contributor Author

Some how the pull request is not synced with my branch. Not sure why?

@alanwaketan
alanwaketan force-pushed the alanwaketan/fix_tpu branch from 959994a to 3e4878f Compare March 13, 2024 01:52
@alanwaketan

alanwaketan commented Mar 13, 2024

Copy link
Copy Markdown
Contributor Author

Some how the pull request is not synced with my branch. Not sure why?

It's resolved. Never mind. Maybe -f is not recognized.

@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Mar 13, 2024

Copy link
Copy Markdown

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.

muellerzr
muellerzr approved these changes Mar 13, 2024
View reviewed changes

@muellerzr muellerzr left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks!

@muellerzr
muellerzr requested a review from amyeroberts March 13, 2024 15:23
amyeroberts
amyeroberts approved these changes Mar 13, 2024
View reviewed changes

@amyeroberts amyeroberts left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for fixing!

@amyeroberts
amyeroberts merged commit b340d90 into huggingface:main Mar 13, 2024
@alanwaketan

alanwaketan commented Mar 13, 2024

Copy link
Copy Markdown
Contributor Author

Thanks for the reviews!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ArthurZucker ArthurZucker ArthurZucker left review comments
+2 more reviewers
@muellerzr muellerzr muellerzr approved these changes
@amyeroberts amyeroberts amyeroberts approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@alanwaketan @HuggingFaceDocBuilderDev @muellerzr @amyeroberts @ArthurZucker