Skip to content

Enable fused foreach Adam compilation #104121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 19 commits into from
Closed

Enable fused foreach Adam compilation #104121

wants to merge 19 commits into from

Conversation

mlazos
Copy link
Contributor

@mlazos mlazos commented Jun 23, 2023 •
edited by pytorch-bot bot
Loading

@mlazos mlazos requested review from albanD and janeyx99 as code owners June 23, 2023 19:50
@pytorch-bot
Copy link

pytorch-bot bot commented Jun 23, 2023 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/104121

Note: Links to docs will display an error until the docs builds have been completed.

✅ 1 Unrelated Failure

As of commit 1cb8777:

BROKEN TRUNK - The following job failed but were present on the merge base 803c144:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@pytorch-bot pytorch-bot bot added the release notes: optimizer Relating to optimizers, torch.optim label Jun 23, 2023
@mlazos
Copy link
Contributor Author

mlazos commented Jun 23, 2023

Currently getting NaNs so I'm investigating that, but wanted to get feedback on the optim changes

janeyx99
janeyx99 reviewed Jun 23, 2023
View reviewed changes
@mlazos mlazos added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 24, 2023
@mlazos
Copy link
Contributor Author

mlazos commented Jun 24, 2023

Currently getting NaNs so I'm investigating that, but wanted to get feedback on the optim changes

Fixed by #104137

@mlazos mlazos requested a review from janeyx99 June 26, 2023 21:18
janeyx99
janeyx99 reviewed Jun 27, 2023
View reviewed changes
Copy link
Contributor

@janeyx99 janeyx99 left a comment

Choose a reason for hiding this comment

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

Instead of using self.compiled and introducing a new "internal-only" flag, which...meh, not good engineering, we can just overload self.capturable. So dynamo can just set capturable to True when compiling.

Then, we're guaranteed that step will follow to the right device. e.g., rn one would need to update line 432 in optimizer.py to check the flag so step could be loaded correctly.

@mlazos mlazos requested a review from janeyx99 June 27, 2023 19:07
@mlazos mlazos force-pushed the mlazos/adam-fused branch from c6c4fe0 to 6f192d4 Compare June 27, 2023 22:10
@pytorch pytorch deleted a comment from linux-foundation-easycla bot Jun 28, 2023
janeyx99
janeyx99 reviewed Jun 29, 2023
View reviewed changes
janeyx99
janeyx99 reviewed Jun 30, 2023
View reviewed changes
Copy link
Contributor

@janeyx99 janeyx99 left a comment

Choose a reason for hiding this comment

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

Did a deeper dive of the changes. I think we should clarify what users should expect when they torch compile vs capturable vs both. I had previously thought of torch.compile be an overriding of capturable…but is there a space where the eager capturable will vary from torch compiled step?

janeyx99
janeyx99 reviewed Jul 5, 2023
View reviewed changes
mlazos and others added 3 commits July 5, 2023 12:23
@mlazos @janeyx99
Update torch/optim/adamw.py
84402fd 
Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>
@mlazos @janeyx99
Update torch/optim/optimizer.py
8769945 
Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>
@mlazos @janeyx99
Update test/inductor/test_compiled_optimizers.py
d8e6c45 
Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>
janeyx99
janeyx99 approved these changes Jul 5, 2023
View reviewed changes
Copy link
Contributor

@janeyx99 janeyx99 left a comment

Choose a reason for hiding this comment

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

🚀 ⭐️

@mlazos
Copy link
Contributor Author

mlazos commented Jul 5, 2023

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: 1 jobs have failed, first few of them are: inductor / cuda11.8-py3.10-gcc7-sm86 / test (inductor_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu)

Details for Dev Infra team Raised by workflow job

@mlazos
Copy link
Contributor Author

mlazos commented Jul 5, 2023

@pytorchbot merge -f "Failures on trunk"

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@huydhn huydhn mentioned this pull request Jul 6, 2023
Closed
@eellison
Copy link
Contributor

eellison commented Jul 7, 2023

I think this is responsible for a 50s increase in TIMM dynamic shapes compilation time 131s → 180s

@janeyx99 janeyx99 mentioned this pull request Jul 10, 2023
Closed
pytorchmergebot pushed a commit that referenced this pull request Jul 10, 2023
@mlazos @pytorchmergebot
Enable running compiled optimizers in CI (#104888)
4063158 
as title

for reference: this is a followup to #104121

Pull Request resolved: #104888
Approved by: https://github.com/janeyx99
@huydhn huydhn mentioned this pull request Jul 12, 2023
Closed
pytorchmergebot pushed a commit that referenced this pull request Jul 14, 2023
@huydhn @pytorchmergebot
Use GitHub REST API to get the merge base commit SHA (#105098)
893983e 
Fixes #104713

### Testing
Manual testing locally using #104121 and confirm that the correct merge base commit is returned [803c144](https://github.com/pytorch/pytorch/commits/1cb87771c1efef32df7009d75bed08249df8ecad) instead of the wrong value provided by `baseRefOid` (de7b6e5).  Here is the JSON output of the GraphQL query for PR info https://paste.sh/TJ-QQWz4#fvE3Y6qoJ8vDkRBZ3vowkZ3m

Pull Request resolved: #105098
Approved by: https://github.com/malfet
@github-actions github-actions bot deleted the mlazos/adam-fused branch January 4, 2025 02:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mikaylagawarecki mikaylagawarecki mikaylagawarecki left review comments

@janeyx99 janeyx99 janeyx99 approved these changes

@albanD albanD Awaiting requested review from albanD albanD is a code owner

Assignees

No one assigned

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request Merged module: dynamo module: inductor release notes: inductor release notes: optimizer Relating to optimizers, torch.optim

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@mlazos @pytorchmergebot @eellison @janeyx99 @mikaylagawarecki