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.

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

Update the tutorial for jax.checkpoint activation offloading policies by add examples #25594

Merged
merged 4 commits into from
Jan 8, 2025
Merged

Update the tutorial for jax.checkpoint activation offloading policies by add examples #25594

merged 4 commits into from
Jan 8, 2025

Conversation

zhenying-liu
Copy link
Contributor

jax.checkpoint supports offloading policies, which are mentioned in tutorials without details.
This CL:

  • creates a subsection for offloading policies.
  • adds two examples to explain activation offloading policies.

mattjj
mattjj approved these changes Dec 18, 2024
View reviewed changes
Copy link
Collaborator

@mattjj mattjj left a comment

Choose a reason for hiding this comment

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

Thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Dec 18, 2024
@zhenying-liu zhenying-liu force-pushed the activation-offloading-doc branch 2 times, most recently from 3a57b72 to 5aa4ab3 Compare December 19, 2024 19:44
@zhenying-liu
Copy link
Contributor Author

I couldn't see the error of "import/copybara — An error happened while migrating the change". Sorry about updating the documentation after it's approved. @jaro-sevcik pointed out the issues in the 2nd code example, so I tested, verified and updated. Please let me know what I need to do to resolve the error. Thanks!

@zhenying-liu zhenying-liu force-pushed the activation-offloading-doc branch from 5aa4ab3 to 522a8fd Compare December 19, 2024 22:49
@zhenying-liu
Copy link
Contributor Author

docs/readthedocs.org:jax has the error message of "the module 'scipy.misc' has no attribute 'face'" which is not related with this CL. The error is addressed in #25728.

@zhenying-liu
Copy link
Contributor Author

Please help rerun the CI as
docs/readthedocs.org:jax — Read the Docs build failed!
had been resolved after merging #25728 and temporarily allowing the warning of deprecation libraries.

@zhenying-liu zhenying-liu changed the title Add examples for jax.checkpoint activation offloading policies Update the tutorial for jax.checkpoint activation offloading policies by add examples Jan 7, 2025
@zhenying-liu
Copy link
Contributor Author

The latest CI results in this CL didn't include the fix for scipy.misc in #25728 merged yesterday.
And it showed the errors for depreciation warnings of SciPy library regarding #25734.
I'm not sure about the Bazel CPU tests. Looks like it couldn't find a container.
Please let me know how I can help pass the CI.

@copybara-service copybara-service bot merged commit 81db321 into jax-ml:main Jan 8, 2025
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yashk2810 yashk2810 yashk2810 approved these changes

@mattjj mattjj mattjj approved these changes

Assignees
No one assigned
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@zhenying-liu @mattjj @yashk2810 @kokoro-team