17 Pull requests merged by 1 person

• Rename convert_to_shlo_fn to convert and jax2obm.convert to convert_polymorphic_fn.

  #1763 merged Mar 29, 2025

• Add placeholder unit test to cover the simple case without tree sharding.

  #1753 merged Mar 28, 2025

• Replace StepMetadata.item_handlers assert with a more helpful ValueError item_handlers is not a dict.

  #1755 merged Mar 28, 2025

• Remove Module class and replace all the places using the Module to the dictionary.

  #1739 merged Mar 27, 2025

• Remove Jax2obm_kwargs flag support_tf_resources now that it is supported by default and is always True.

  #1746 merged Mar 26, 2025

• #v1 Add metadata free functions.

  #1691 merged Mar 26, 2025

• Internal

  #1724 merged Mar 26, 2025

• Introduce Placeholder to PyTreeCheckpointHandler. This allows PyTrees to be partially restored by specifying which nodes should be skipped using a Placeholder value.

  #1732 merged Mar 26, 2025

• Wrap synchronous checkpoint handler calls in composite checkpoint handler with CommitFutureAwaitingContractedSignals to not break save.

  #1736 merged Mar 26, 2025

• set jax_enable_x64 to True to test int64

  #1748 merged Mar 26, 2025

• Add some logs for debugging awaitable signals.

  #1744 merged Mar 26, 2025

• Remove version argument in export manager constructor.

  #1745 merged Mar 26, 2025

• Broadcast interval time-based checkpoint decision across hosts in CheckpointManager from primary host.

  #1741 merged Mar 25, 2025

• Remove version argument in export manager constructor.

  #1731 merged Mar 24, 2025

• Remove all set directory signals in jax distributed client after async save is complete.

  #1733 merged Mar 24, 2025

• Disable async directory creation in composite checkpoint handler if checkpointer is not asynchronous.

  #1734 merged Mar 24, 2025

• Add is_leaf parameter to tree_difference.

  #1702 merged Mar 24, 2025

13 Pull requests opened by 1 person

• Use `ocp.future.CommitFutureAwaitingContractedSignals` in custom async checkpoint handlers and type handlers to support async directory creation.

  #1737 opened Mar 24, 2025

• Use ocp.futures.CommitFutureAwaitingContractedSignals to support async directory creation feature for Gemax Prod

  #1738 opened Mar 24, 2025

• Patch for testing b/404645078, hardcoding single_host_load_and_broadcast=true.

  #1740 opened Mar 25, 2025

• Adding preservation policy

  #1747 opened Mar 26, 2025

• #v1 Add CheckpointableHandlerRegistry and handler resolution functions.

  #1749 opened Mar 26, 2025

• Patch for `use_replica_parallel=False`.

  #1751 opened Mar 26, 2025

• #v1 Add `CompositeHandler`, similar to the current `CompositeCheckpointHandler`, but different in the sense that it will not be visible to the user and is solely a helper layer for save/load functions.

  #1752 opened Mar 26, 2025

• Add utility function for filtering a PyTree to match the structure of a target PyTree.

  #1754 opened Mar 27, 2025

• Updated checkpoint loading to always load the latest valid checkpoint, instead of choosing any checkpoint without preference. This significantly reduces compute as it minimizes re-computation of already computed steps.

  #1759 opened Mar 28, 2025

• Catching very infrequent error in train-logs writing which crashes train job.

  #1760 opened Mar 28, 2025

• internal

  #1761 opened Mar 28, 2025

• #v1 Implement `save/load_checkpointables` in terms of V0 APIs. Resolve handlers for checkpointables using V1 APIs.

  #1762 opened Mar 28, 2025

• Add `tree_trim` function for filtering a PyTree by another's structure.

  #1764 opened Mar 28, 2025

1 Issue closed by 1 person

• Weights of Restored model after checkpointing do not give same loss as model before saving

  #1402 closed Mar 28, 2025

2 Unresolved conversations

Sometimes conversations happen on old items that aren’t yet closed. Here is a list of all the Issues and Pull Requests with unresolved conversations.

• Incorrect / inconvenient type annotations inside orbax.checkpoint.checkpoint_args.register_with_handler

  #1718 commented on Mar 25, 2025 • 0 new comments

• Make the `create_directories_asynchronously` option default to `True` if jax distributed client is initialized.

  #1630 commented on Mar 28, 2025 • 0 new comments