Restoring to CPU

[lachinov]

Hi,

I'm using orbax checkpoint version '0.5.20'
After moving to the new API, I got an issue similar to #646 , but with a twist. I don't know the structure of the PyTree beforehand.
The code snippet for restoration looks as follows and works fine when the device configuration is the same as during the training. I want to read the saved data to the CPU and from there apply my sharding.

checkpoint = checkpoint_manager.restore(step=best_step,
                                            args=orbax.checkpoint.args.Composite(
                                                configuration=orbax.checkpoint.args.PyTreeRestore(),
                                                metric_values=orbax.checkpoint.args.PyTreeRestore(),
                                                hparams=orbax.checkpoint.args.JsonRestore(),
                                            ),
                                            )

It would be possible if I knew configuration signature.
However, configuration tree structure is learnable and dynamic, so I don't know its signature beforehand.

Is there any solution besides changing the saving logic?

[cpgaffney1]

You should be able to do this:

with CheckpointManager(directory, options, item_handlers={'configuration': ocp.PyTreeCheckpointHandler(), ...}) as mngr:
  step = mngr.latest_step()
  metadatas = mngr.item_metadata(step)
  configuration_metadata = metadatas.configuration
  # Map configuration_metadata, which is a tree structure, to the shardings you want to restore with.
  mngr.restore(...)

The item_handlers is require up front to use item_metadata() method. In the (near-ish) future, this configuration will not be required.

Also note that you would be able to specify restore_args with shardings so that you can get orbax to load directly onto devices, rather than first onto CPUs and then manually apply shardings. But you can do it that way if you want.

[lachinov]

Thanks! That's exactly what I needed.