Orbax Loading and Sharding Support feature #21903
Conversation
Summary of ChangesHello @amitsrivastava78, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances Keras's checkpoint loading capabilities by introducing a dedicated Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces significant new functionality for Orbax checkpoint loading and sharding support. The changes include a new Model.load() method, helper functions for checkpoint discovery, and a comprehensive suite of new tests. The implementation is generally solid and well-tested. My feedback focuses on improving code maintainability by reducing duplication and enhancing error handling in the new test cases.
amitsrivastava78
force-pushed
the
orbax-pending-features
branch
from
ee32491 to
10370f5
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #21903 +/- ##
==========================================
- Coverage 76.30% 76.30% -0.01%
==========================================
Files 580 580
Lines 60029 60112 +83
Branches 9432 9450 +18
==========================================
+ Hits 45803 45866 +63
- Misses 11750 11758 +8
- Partials 2476 2488 +12
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
amitsrivastava78
force-pushed
the
orbax-pending-features
branch
from
10370f5 to
325a03f
Compare
amitsrivastava78
force-pushed
the
orbax-pending-features
branch
from
325a03f to
43d45d0
Compare
- Remove complex JAX abstract pytree logic that was causing 'free(): invalid pointer' errors - Use preservation mode for all backends to avoid state structure mismatches - This prevents memory corruption when loading checkpoints with different optimizer states
- Replace bare 'except:' with specific 'except (ImportError, AttributeError):' for distribution import patterns - This improves error handling by only catching expected exceptions
- Extract duplicated tensor conversion logic into _to_numpy() helper method - Replace duplicated code blocks in optimizer and metrics variable comparisons - Improves maintainability and reduces code duplication
amitsrivastava78
force-pushed
the
orbax-pending-features
branch
from
c592a27 to
d8a86e8
Compare
| try: | ||
| import orbax.checkpoint as ocp | ||
|
|
||
| return ocp.multihost.is_initialized() |
| return False | ||
|
|
||
| try: | ||
| import orbax.checkpoint as ocp |
There was a problem hiding this comment.
I don't quite understand the reason for all the ad-hoc imports. At least, could we factor out that repeated pattern to reduce code ugliness?
| self._sync_processes(f"checkpoint_save_start_{step}") | ||
|
|
||
| # --- Save Logic (V1 API) --- | ||
| if self.verbose > 0 and self._is_primary_host(): |
There was a problem hiding this comment.
Oh you shouldn't do this, Orbax expects all APIs to be called from every host. Orbax also handles barriers internally.
| # save_pytree is called | ||
| if self.verbose > 0: | ||
| # Synchronize before saving to ensure consistency | ||
| self._sync_processes(f"checkpoint_save_start_{step}") |
There was a problem hiding this comment.
You can do this if you want but it shouldn't affect correctness of the underlying save.
| h5py = None | ||
|
|
||
|
|
||
| def _is_orbax_checkpoint(filepath): |
There was a problem hiding this comment.
Please don't look at internal files, it's a bit of a brittle behavior.
You can just use this function: https://source.corp.google.com/piper///depot/google3/third_party/py/orbax/checkpoint/experimental/v1/__init__.py;l=27
| return False | ||
|
|
||
|
|
||
| def _find_latest_orbax_checkpoint(checkpoint_dir): |
There was a problem hiding this comment.
checkpointer.latest can replace this impl
| metrics state) from an Orbax checkpoint directory. The checkpoint | ||
| directory should contain subdirectories named with step numbers. | ||
| If the filepath points to a checkpoint directory, it will load the |
There was a problem hiding this comment.
To be pedantic, I think it helps users if we are very strict about terminology. The user may ask, "what is the difference between a checkpoint directory and a specific step directory? I thought the step directory IS a checkpoint directory?" The solution in my view is to distinguish between the root directory, which is a container for a list of checkpoint/step directories.
5 participants
Support