-
Notifications
You must be signed in to change notification settings - Fork 233
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
Feature request: Use Orbax for checkpointing. #80
Comments
I think the current streaming checkpointer is quite good and I'm not planning to switch to orbax. Specifically, the streaming checkpointer can save and load checkpoints with minimal memory and temporary disk usage, which means that you can checkpoint models that does not fit in the memory or local disk of a single machine. Also as a personally preference, I want to stay away from complicated Google libraries as much as possible. |
obrax do not support load streaming and sharding data or array across devices with pjit so I think the current checkpointing method that is being used right now is a smart move :) |
wdyt about changing the current checkpointing to use orbax?
https://github.com/google/orbax
The text was updated successfully, but these errors were encountered: