orbax

orbax copied to clipboard

Add Emergency Checkpoint Manager Benchmark

copybara-service[bot]

Support Colocated Python Checkpointing

copybara-service[bot]

Use `SingleReplicaArrayHandler` for restoring `jax.Array` in emergency checkpoint manager.

1

Use `SingleReplicaArrayHandler` for restoring `jax.Array` in emergency checkpoint manager.

copybara-service[bot]

Fix listing perf for gcs by completely avoiding the extra list calls

1

Fix listing perf for gcs by completely avoiding the extra list calls

copybara-service[bot]

Internal change

copybara-service[bot]

Internal change

1

Internal change

copybara-service[bot]

Allow reading locally empty arrays in multi-controller jax

copybara-service[bot]

Fix readthedocs build file. Add the oex_orchestration.proto to the proto compile function.

1

Fix readthedocs build file. Add the oex_orchestration.proto to the proto compile function.

copybara-service[bot]

Internal.

copybara-service[bot]

[replica-parallel] Configurable limits for replica-parallel checkpoint saving

1

**Adds parameters `min_slice_bytes_for_replica_parallel` and `max_replicas_for_replica_parallel` to `ArrayHandler` to allow users to configure replica-parallel checkpoint saving.** If `use_replica_parallel` is set to true, saving will be parallelized over at most `max_replicas_for_replica_parallel` different...

rao-ashish