Skip to content
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.

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

schedule-free optimier: ensure it's possible to donate both the state and the params #1059

Merged

Conversation

enolan
Copy link
Contributor

@enolan enolan commented Sep 13, 2024

Prior to this commit, with a schedule-free optimizer, if you tried to donate an entire train state, including both the parameters and a freshly initialized optimizer state, you'd get an "INVALID_ARGUMENT: Attempt to donate the same buffer twice in Execute()" error, because z was the same array as params. This commit fixes the issue and adds a test.

@vroulet
Copy link
Collaborator

vroulet commented Sep 14, 2024

Hello @enolan,

Thank you for catching that!
Could you wait for #1060 to be merged ?
I can ping you back when it's merged into head, you'll merge your PR with head, then I'll review.

@vroulet vroulet self-assigned this Sep 14, 2024
@enolan
Copy link
Contributor Author

enolan commented Sep 14, 2024

sure no problem

@vroulet
Copy link
Collaborator

vroulet commented Sep 23, 2024

Hello @enolan,
Your PR should be good to go. Just sync with head, check if the problem persists without your change (i.e. whether #1060 changed the behavior), then upload, and I'll approve.
Thanks again for looking into this!

…arams

Prior to this commit, with a schedule-free optimizer, if you tried to donate
an entire train state, including both the parameters and a freshly initialized
optimizer state, you'd get an "INVALID_ARGUMENT: Attempt to donate the same
buffer twice in Execute()" error, because z was the same array as params. This
commit fixes the issue and adds a test.
@enolan
Copy link
Contributor Author

enolan commented Sep 29, 2024

We should be good. I rebased and checked, the test still fails without the copy and passes with it.

Copy link
Collaborator

@vroulet vroulet left a comment

Choose a reason for hiding this comment

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

Thank you very much for looking into that.
The comment and the test are great

@copybara-service copybara-service bot merged commit dd1daab into google-deepmind:main Oct 2, 2024
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants