Optimized Checkpointing with Tensorstore#
Orbax relies on Tensorstore to store individual arrays in a checkpoint. Tensorstore provides efficient, scalable library for reading and writing arrays.
Until recently, however, our use of Tensorstore came with a few drawbacks. Chief among them was the fact that every parameter in a training state would be saved as a separate directory. This approach can be quite performant, even for models with hundreds of billions of parameters, provided that model layers are stacked. Otherwise, hundreds or thousands of directories may be created in the checkpoint.
This fact can lead to very slow restore times, which is undesirable in and of itself, but is particularly painful for jobs that may be preempted frequently and need to restart, for example.
While it is slightly less of a concern at save time, since writes to disk can happen asynchronously, the synchronous portion of the save can still be slow as many directories are created.
Additionally, if individual parameters are small, storage may be wasted on filesystems with minimum file sizes.
Towards an Improved Checkpoint Format#
The new, optimized checkpoint format provided by Orbax is backed by Tensorstore’s OCDBT driver.
For practical purposes, this means that we will no longer store one parameter per directory, but will aggregate many parameters into a smaller set of large files.
Empirically, we have observed substantial speed-ups in both save and restore when using the new format.
Save Performance (sec)#
Restore Performance (sec)#
Checkpoint Format#
Concretely, what does the new checkpoint format look like in comparison to the old?
Old Format#
f = """
path/to/my/checkpoint/dir/
0/
state/
layer0.param0/
.zarray
0.0
0.1
1.0
1.1
layer1.param0/
.zarray
0.0
...
<another_item>/
...
1/
...
2/
...
Note: in this case, `0.0`, `0.1`, etc. provides an indication of how the array
was sharded when originally saved.
"""
New Format#
f = """
path/to/my/checkpoint/dir/
0/
state/
checkpoint # legacy msgpack file, stores tree structure
tree_metadata # (maybe) new proto file, stores tree structure
d/ # array data stored here
012b2c6e5c9d2a16c240a59d5f0f35c0
056e0816bdc5496a86251e58a0ec202b
...
manifest.0000000000000001
...
manifest.ocdbt
<another_item>/
...
1/
...
2/
...
"""
Enabling the new format#
import jax
import tempfile
import subprocess
import os
from etils import epath
import orbax.checkpoint as ocp
# Initialize PyTreeCheckpointHandler with `use_ocdbt=True`.
# This option already defaults to True, so it's optional to pass it in.
ckptr = ocp.Checkpointer(ocp.PyTreeCheckpointHandler(use_ocdbt=True))
Additional Notes#
All checkpoints previously produced by Orbax in the old format will still be
readable when the new format is enabled. However, if a checkpoint is produced in the new format, it cannot be read if use_ocdbt is disabled.
Custom Chunk Sizes#
Orbax Zarr3, a multidimensional array storage format, offers customizable chunk sizes in bytes for optimal memory management. The default chunk size, which corresponds one-to-one with the array shard size, can cause out-of-memory errors when reading on hosts with different sharding layouts. For example, this can often arise when arrays are saved with a fully-sharded sharding, but loaded with a fully-replicated sharding. To prevent this, set chunk_byte_size smaller than or equal to the anticipated read size. Anything above 1MB generally won’t affect impact performance. Consider following example:
# setup checkpoint data
array_len = 8 * 1024
key = jax.random.PRNGKey(0)
key, subkey = jax.random.split(key)
pytree = {
'a': jax.random.normal(subkey, (array_len, ), dtype=jax.numpy.float32), # 32KB
'b': jax.random.normal(subkey, (array_len * 2, ), dtype=jax.numpy.float32), # 64KB
}
# create save_args to customize the chunk_byte_size
save_args = jax.tree_util.tree_map(
lambda x: ocp.SaveArgs(
chunk_byte_size=
1024, # 1KB
),
pytree,
)
temp_dir = tempfile.TemporaryDirectory()
mgr = ocp.CheckpointManager(epath.Path(temp_dir.name),
item_handlers=ocp.PyTreeCheckpointHandler(use_zarr3=True)) # make sure zarr3 is enabled
mgr.save(
0,
args=ocp.args.PyTreeSave(
pytree,
save_args=save_args,
),
)
mgr.close()
WARNING:absl:Setting the target_byte_size too small could reduce performance.
WARNING:absl:Setting the target_byte_size too small could reduce performance.
Customizing Data File Size#
To improve file I/O parallelism when working with large files on remote storages like GCS, use the PyTreeSaveArgs.ocdbt_target_data_file_size parameter to control the size of output files.
BEFORE#
def print_directory_file_size(dir: epath.Path) -> None:
print(f"dir={dir}:")
for f in data_dir.iterdir():
if f.is_file():
print(f"file={f.name}, size={f.stat().length}")
# continue from above example, examine the data file sizes
data_dir = epath.Path(temp_dir.name) / '0'/ 'default'/ 'ocdbt.process_0'/ 'd'
print_directory_file_size(data_dir)
dir=/tmp/tmpx9gpqfny/0/default/ocdbt.process_0/d:
file=ba44b04c15b77068fbcfd85e9dfe8406, size=304
file=3c9d09486f92edb85c1c374e971eec01, size=33312
file=24054ec2a0ab050ce50de955812afaae, size=285
file=b77f6451e7bb3bf4da4d7824cf4e01ca, size=66351
AFTER#
temp_dir = tempfile.TemporaryDirectory()
mgr = ocp.CheckpointManager(temp_dir.name,
item_handlers=ocp.PyTreeCheckpointHandler(use_zarr3=True))
mgr.save(
0,
args=ocp.args.PyTreeSave(
pytree,
save_args=save_args,
# 400 MiB, should be much larger than chunk_byte_size
ocdbt_target_data_file_size=400 * 2**20,
),
)
mgr.close()
data_dir = epath.Path(temp_dir.name) / '0'/ 'default'/ 'ocdbt.process_0'/ 'd'
print_directory_file_size(data_dir)
WARNING:absl:Setting the target_byte_size too small could reduce performance.
ERROR:absl:[process=0][thread=MainThread][step=0][wait_until_finished] Save Finalize thread (save_finalize) failed.
Traceback (most recent call last):
File "/home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/checkpoint/checkpoint_manager.py", line 2057, in wait_until_finished
self._finalize_thread.get_not_none().join()
File "/home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/checkpoint/checkpoint_manager.py", line 157, in join
super().join(*args, **kwargs)
File "/home/docs/.asdf/installs/python/3.12.13/lib/python3.12/threading.py", line 1149, in join
self._wait_for_tstate_lock()
File "/home/docs/.asdf/installs/python/3.12.13/lib/python3.12/threading.py", line 1169, in _wait_for_tstate_lock
if lock.acquire(block, timeout):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt
WARNING:absl:Setting the target_byte_size too small could reduce performance.
---------------------------------------------------------------------------
KeyboardInterrupt Traceback (most recent call last)
Cell In[8], line 15
2 mgr = ocp.CheckpointManager(temp_dir.name,
3 item_handlers=ocp.PyTreeCheckpointHandler(use_zarr3=True))
5 mgr.save(
6 0,
7 args=ocp.args.PyTreeSave(
(...) 12 ),
13 )
---> 15 mgr.close()
17 data_dir = epath.Path(temp_dir.name) / '0'/ 'default'/ 'ocdbt.process_0'/ 'd'
18 print_directory_file_size(data_dir)
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/checkpoint/checkpoint_manager.py:2184, in CheckpointManager.close(self)
2182 def close(self):
2183 """Waits for outstanding operations to finish and closes internal objects."""
-> 2184 self.wait_until_finished()
2185 self._checkpointer.close()
2186 # Call after checkpointer.close().
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/checkpoint/checkpoint_manager.py:2057, in CheckpointManager.wait_until_finished(self)
2046 logging.info(
2047 '[process=%s][thread=%s][step=%s][wait_until_finished] Waiting for'
2048 ' Save Finalize thread (%s) to complete.',
(...) 2052 finalize_thread_name,
2053 )
2054 # Let all threads join and wait for the finalize thread to complete.
2055 # Don't call join() with a lock otherwise we will end up serializing the
2056 # access to the finalize thread.
-> 2057 self._finalize_thread.get_not_none().join()
2058 logging.info(
2059 '[process=%s][thread=%s][step=%s][wait_until_finished] Done'
2060 ' waiting for Save Finalize thread (%s) running at step=%d.',
(...) 2065 step,
2066 )
2067 except BaseException: # pylint:disable=broad-exception-caught
File ~/checkouts/readthedocs.org/user_builds/orbax/envs/latest/lib/python3.12/site-packages/orbax/checkpoint/checkpoint_manager.py:157, in _FinalizeThread.join(self, *args, **kwargs)
156 def join(self, *args, **kwargs):
--> 157 super().join(*args, **kwargs)
159 # Keep track of whether the calling thread should raise the exception.
160 if not hasattr(self._thread_local, 'should_raise'):
File ~/.asdf/installs/python/3.12.13/lib/python3.12/threading.py:1149, in Thread.join(self, timeout)
1146 raise RuntimeError("cannot join current thread")
1148 if timeout is None:
-> 1149 self._wait_for_tstate_lock()
1150 else:
1151 # the behavior of a negative timeout isn't documented, but
1152 # historically .join(timeout=x) for x<0 has acted as if timeout=0
1153 self._wait_for_tstate_lock(timeout=max(timeout, 0))
File ~/.asdf/installs/python/3.12.13/lib/python3.12/threading.py:1169, in Thread._wait_for_tstate_lock(self, block, timeout)
1166 return
1168 try:
-> 1169 if lock.acquire(block, timeout):
1170 lock.release()
1171 self._stop()
KeyboardInterrupt: