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

fix: Coherent dtypes of updates with and without MultiSteps #1122

Merged
merged 2 commits into from
Nov 5, 2024

Conversation

hlzl
Copy link
Contributor

@hlzl hlzl commented Oct 28, 2024

Following #1039 and discussions with @vroulet a proposed fix such that updates with and without MultiSteps have the same dtypes independent of how acc_grads and the inner_opt_state were initialized.

hlzl added 2 commits October 28, 2024 17:38
Test cases check different combinations of parameter and update dtypes.
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.

Very nice, thank you @hlzl !
By curiosity did you check the speed before and after your change?
I think it should be the same but would be good to check just to be sure.

In any case, I'll approve this, thank you for all the effort!

@copybara-service copybara-service bot merged commit c4fd723 into google-deepmind:main Nov 5, 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