msgpack_save

Contents

msgpack_save#

class braintools.file.msgpack_save(filename, target, overwrite=True, async_manager=None, verbose=False)[source]#

Save a checkpoint of the model. Suitable for single-host using the msgpack library.

This function is rewritten from the Flax APIs (google/flax).

In this method, every JAX process saves the checkpoint on its own. Do not use it if you have multiple processes and you intend for them to save data to a common directory (e.g., a GCloud bucket). To save multi-process checkpoints to a shared storage or to save GlobalDeviceArray`s, use `multiprocess_save() instead.

Pre-emption safe by writing to a unique temporary file before a final rename. When async_manager is given, the file write runs in the background; the target is serialized synchronously before returning so the on-disk snapshot is consistent even if the model keeps training. Wait for the write with async_manager.wait_previous_save().

Parameters:

  • filename (str | PathLike) – path to store the checkpoint file in.

  • target (PyTree) – serializable object.

  • overwrite (bool) – overwrite an existing checkpoint at filename (default: True). When False, an existing file raises InvalidCheckpointPath.

  • async_manager (AsyncManager | None) – if defined, the file write will run without blocking the main thread. Only works for single host. Note that an ongoing save will still block subsequent saves, to make sure overwrite/keep logic works correctly.

  • verbose (bool) – Whether to print progress information (default: False).

Returns:

out – Filename of the saved checkpoint.

Return type:

str

Contents