orbax

orbax copied to clipboard

internal change

1

internal change

copybara-service[bot]

Internal.

copybara-service[bot]

Internal change.

copybara-service[bot]

Create JaxModule Protocol class. So we can separate the interface and Implementation.

copybara-service[bot]

SAX implementation of gemini-phase2 MMVisionModel.

copybara-service[bot]

Add TypeHandlers to regeister default SaveArg/RestoreArgs. This allows users to save with `Checkpointer.save(..., save_arg=None)` (or restore with `restore_arg=None`), and have the argument automatically be converted to the default `SaveArg`/`RestoreArg` for...

copybara-service[bot]

Pass `int` to grain seed

copybara-service[bot]

Incorrect null check in pytree_checkpoint_handler.py

4

`orbax/checkpoint/pytree_checkpoint_handler.py:661` has the following check: `if not item` It most likely should be `if item is None`, as otherwise this check will raise an error when item is an array...

hr0nix

Sharded loading performance question

10

I'm benchmarking loading a 65B sharded transformer model on multiple GPUs on the same host. The checkpoint itself is not sharded, but when the model is being loaded, a correct...

hr0nix

Add retention policy for pytreecheckpointhandler metrics

copybara-service[bot]