Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Orbax Loading and Sharding Support feature#21903

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 ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Open
amitsrivastava78 wants to merge18 commits intokeras-team:master
base:master
Choose a base branch
Loading
fromamitsrivastava78:orbax-pending-features

Conversation

@amitsrivastava78
Copy link
Collaborator

Support

  1. New loading function for the orbax checkpointing
  2. Support sharding feature

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello@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 dedicatedModel.load() method for Orbax checkpoints. This new functionality intelligently handles checkpoint directories, automatically identifying and loading the latest state. Crucially, it incorporates advanced sharding support for the JAX backend, allowing models to be loaded and resharded dynamically to match the current distribution configuration. Additionally, theModel.export() method has been refined by removing the LiteRT export option, focusing on core export functionalities.

Highlights

  • New Model.load() method for Orbax Checkpoints: A newload method has been added to theModel class, enabling direct loading of Orbax checkpoints. This method can automatically detect the latest checkpoint in a directory or load a specific step.
  • Sharding Support for JAX Backend: TheModel.load() method now supports automatic resharding of data when loading Orbax checkpoints on the JAX backend, adapting to the current distribution strategy.
  • Refined Model.export() Functionality: TheModel.export() method has been updated to remove support for LiteRT export, streamlining the available export formats.
  • Comprehensive Testing for Loading and Sharding: Extensive new tests have been added toorbax_checkpoint_test.py to validate the newModel.load() method, including synchronous/asynchronous saving, layout preservation, and JAX-specific resharding and distributed checkpoint structure verification.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on ourdocumentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either/gemini <command> or@gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

FeatureCommandDescription
Code Review/gemini reviewPerforms a code review for the current pull request in its current state.
Pull Request Summary/gemini summaryProvides a summary of the current pull request in its current state.
Comment@gemini-code-assistResponds in comments when explicitly tagged, both in pull request comments and review comments.
Help/gemini helpDisplays a list of available commands.

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.gemini/ folder in the base of the repository. Detailed instructions can be foundhere.

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 uphere.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with theGemini Code Assist IDE Extension.

Footnotes

  1. Review thePrivacy Notices,Generative AI Prohibited Use Policy,Terms of Service, and learn how to configure Gemini Code Assist in GitHubhere. Gemini can make mistakes, so double check it anduse code with caution.↩

Copy link
Contributor

@gemini-code-assistgemini-code-assistbot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces significant new functionality for Orbax checkpoint loading and sharding support. The changes include a newModel.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.

@codecov-commenter
Copy link

codecov-commenter commentedDec 8, 2025•
edited
Loading

Codecov Report

❌ Patch coverage is68.36735% with62 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.61%. Comparing base (f0a48a6) to head (1f28e74).
⚠️ Report is 17 commits behind head on master.

Files with missing linesPatch %Lines
keras/src/saving/saving_api.py64.34%28 Missing and 18 partials⚠️
keras/src/callbacks/orbax_checkpoint.py75.51%9 Missing and 3 partials⚠️
keras/src/saving/saving_lib.py0.00%0 Missing and 2 partials⚠️
keras/src/utils/module_utils.py77.77%1 Missing and 1 partial⚠️
Additional details and impacted files
@@            Coverage Diff             @@##           master   #21903      +/-   ##==========================================+ Coverage   76.30%   82.61%   +6.31%==========================================  Files         580      588       +8       Lines       60029    61416    +1387       Branches     9432     9653     +221     ==========================================+ Hits        45803    50740    +4937+ Misses      11750     8168    -3582- Partials     2476     2508      +32
FlagCoverage Δ
keras82.43% <67.85%> (+6.26%)⬆️
keras-jax61.68% <67.34%> (-0.45%)⬇️
keras-numpy56.76% <10.71%> (-0.57%)⬇️
keras-openvino37.10% <7.65%> (+2.80%)⬆️
keras-tensorflow63.84% <65.81%> (?)
keras-torch62.57% <65.81%> (-0.65%)⬇️

Flags with carried forward coverage won't be shown.Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report?Share it here.

🚀 New features to boost your workflow:
  • ❄️Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

- 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
- Add multi-host support using orbax.checkpoint.multihost APIs- Remove manual sync calls around save operations- Use proper Orbax v1 APIs instead of brittle file inspection- Fix 80-column line length violations in test files- Ensure cross-backend compatibility with appropriate test skipping- Clarify checkpoint directory terminology in documentation
Copy link
Collaborator

@hertschuhhertschuh left a comment

Choose a reason for hiding this comment

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

Thanks for the PR!

Can you tighten theorbax_checkpoint_test.py file. It's extremely long and hard to follow. I think:

  • a lot fewer tests could cover basically the same
  • some parameterized tests could minimize code duplication
  • the verification blocks could be much shorter usingkeras.tree andself.assertAllClose, I gave some examples.

Copy link

@orbax-devorbax-dev left a comment

Choose a reason for hiding this comment

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

LGTM, it won't let me approve with comments.

Copy link
Collaborator

@hertschuhhertschuh left a comment

Choose a reason for hiding this comment

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

A couple of complications I didn't think about:

- Remove redundant try/except blocks in favor of LazyModule error handling- Use ocp.multihost directly from LazyModule instead of custom function- Remove unnecessary dummy apply_gradients in Orbax loading (confirmed required)- Update sync key to be process-safe- Remove unused test case for asset support- Improve LazyModule to expose multihost from parent orbax.checkpoint module
- Remove manual directory cleanup in _save_checkpoint that interfered with Orbax preservation policies- Simplify preservation policy setup to use LatestN directly instead of AnyPreservationPolicy wrapper- Update asset directory structure to use checkpoint_dir/assets/step/ format- Add comprehensive asset saving/loading tests for both sync and async modes- Make test_save_freq_epoch more robust by checking for numeric checkpoint names rather than specific epoch- Fix asset loading to handle new directory structure in saving_api.pyAll Orbax checkpoint tests now pass on both JAX and TensorFlow backends.
- Remove manual directory cleanup in _save_checkpoint that interfered with Orbax preservation policies- Simplify preservation policy setup to use LatestN directly instead of AnyPreservationPolicy wrapper- Update asset directory structure to use checkpoint_dir/assets/step/ format- Add comprehensive asset saving/loading tests for both sync and async modes- Make test_save_freq_epoch more robust by checking for numeric checkpoint names rather than specific epoch- Fix asset loading to handle new directory structure in saving_api.pyAll Orbax checkpoint tests now pass on both JAX and TensorFlow backends.
Copy link
Collaborator

@hertschuhhertschuh left a comment

Choose a reason for hiding this comment

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

I didn't see where the hook is to supportmodel.load_weights. Did I miss something?

- Save custom layer assets (binary data, strings, arrays) directly in the pytree as base64-encoded strings for Orbax compatibility- Remove separate asset file saving to eliminate synchronization races- Update loading to extract assets from pytree and decode back to original types- Modify tests to verify asset loading without directory checks- Ensures atomic saves with proper preservation policy handling
…points- Add asset loading logic to saving_api.load_weights for Orbax checkpoints- Include assets in weights-only Orbax checkpoints- Fix layer naming in MockLayerWithAssets test to avoid conflicts- Add test for load_weights with assets from Orbax checkpoints
@amitsrivastava78
Copy link
CollaboratorAuthor

I didn't see where the hook is to supportmodel.load_weights. Did I miss something?

Added asset loading logic to saving_api.load_weights for Orbax checkpoints
Included assets in weights-only Orbax checkpoints
Added test for load_weights with assets from Orbax checkpoints

Sign up for freeto join this conversation on GitHub. Already have an account?Sign in to comment

Reviewers

@hertschuhhertschuhhertschuh left review comments

+2 more reviewers

@gemini-code-assistgemini-code-assist[bot]gemini-code-assist[bot] left review comments

@orbax-devorbax-devorbax-dev left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@gbanedgbaned

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@amitsrivastava78@codecov-commenter@hertschuh@orbax-dev@kokoro-team@gbaned
[8]ページ先頭
©2009-2025 Movatter.jp