Error with PytreeCheckpointHandler

[svarunid]

I'm trying to save the model and the optimizer state. I use PytreeCheckpointer with a checkpoint manager for saving. However, I get the following error in the process. The error originates from merge_ocdbt_per_process_files method.

Here's my code:

mngr.save(
    i * (e + 1), # batch * epoch + 1
    {
        "model": model,
        "opt_state": opt_state,
    }
)

Below I attach error message:

  File "D:\Projects\Data Science\machine_translation\env\Lib\site-packages\orbax\checkpoint\checkpoint_manager.py", line 515, in save
    self._checkpointers[k].save(item_dir, item, **kwargs)
  File "D:\Projects\Data Science\machine_translation\env\Lib\site-packages\orbax\checkpoint\checkpointer.py", line 150, in save
    self._handler.finalize(tmpdir)
  File "D:\Projects\Data Science\machine_translation\env\Lib\site-packages\orbax\checkpoint\pytree_checkpoint_handler.py", line 1411, in finalize  
    type_handlers.merge_ocdbt_per_process_files(directory)
  File "D:\Projects\Data Science\machine_translation\env\Lib\site-packages\orbax\checkpoint\type_handlers.py", line 744, in merge_ocdbt_per_process_files
    asyncio.run(open_and_copy())
  File "D:\Projects\Data Science\machine_translation\env\Lib\asyncio\runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "D:\Projects\Data Science\machine_translation\env\Lib\asyncio\runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\Projects\Data Science\machine_translation\env\Lib\asyncio\base_events.py", line 653, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "D:\Projects\Data Science\machine_translation\env\Lib\site-packages\orbax\checkpoint\type_handlers.py", line 742, in open_and_copy
    await asyncio.gather(*copy_ops)
  File "D:\Projects\Data Science\machine_translation\env\Lib\asyncio\tasks.py", line 684, in _wrap_awaitable
    return (yield from awaitable.__await__())
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: CopyRange not supported

[niketkumar]

Which version of Orbax are you using? Can you please try it with the latest version?

What storage are you using to save the checkpoints to? It seems you are running it on Windows. It could be the front/back slashes in the storage path. Can you please try it on Linux?

If the above doesn't work then to get unblocked, can you please try disabling ocdbt_merge with PytreeCheckpointer(..., ocdbt_merge=False, ...)?

[svarunid]

I'm using the latest version (orbax-checkpoint-0.4.8) of orbax-checkpoint. Yes, I'm on windows. I didn't quite understand what you mean by the storage I'm using. I'm using tensorstore. Is this what you were asking?

About using the forward/backward slashes, I use the epath from etils library to construct the path object and then resolve it to platform specific path while initialising the checkpoint manager. Hence, I guess it shouldn't be a problem.

[niketkumar]

Thanks for confirming version and storage (windows, linux, Google Cloud Storage)!

Can you please share the code snippet where CheckpointManager is getting created and if possible how model and opt_state objects are created?

[svarunid]

Sure, however, it would be a large number of lines of code. So, let me just provide a GitHub link for the code snippet I use.
Model is just a PyTree structure with custom nodes in it registered using dataclasses.
And here's the link to the function optim

Here's how I initialize the opt_state and the CheckpointManager:

lr = 0.001
lr = optax.warmup_cosine_decay_schedule(lr, 0.1, 200, 2000, 0.0001)
optimizer = optax.inject_hyperparams(optax.adam)(learning_rate=lr)
del lr

opt_state, step = optim(
    model, optimizer, loss, in_axes=(None, 0, 0, 0, 0, 0), out_axes=0
)

freq = 25
checkpoint_path = epath.Path("path_to_the_directory")
if checkpoint_path.exists() and should_overwrite:
    checkpoint_path.rmtree()
    checkpoint_path.mkdir(parents=True, exist_ok=True)
options = ocp.CheckpointManagerOptions(
    max_to_keep=3,
    save_interval_steps=freq,
)
mngr = ocp.CheckpointManager(
    checkpoint_path.resolve(),
    {
        "model": ocp.AsyncCheckpointer(
            ocp.PyTreeCheckpointHandler(ocdbt_merge=False) # ocdbt_merge=False doesn't produce errors
        ),
        "opt_state": ocp.AsyncCheckpointer(
            ocp.PyTreeCheckpointHandler(ocdbt_merge=False)
        ),
    },
    options,
)

[niketkumar]

Thanks Arun, for sharing the details and confirming that ocdbt_merge=False has unblocked you.

I will try to reproduce the issue with the details.

[niketkumar]

Windows support is limited as of now. We will take it up in the future.

[svarunid]

Thanks. I've switched to Linux and everything works fine!