Saving checkpoint by parts

[jaro-sevcik]

Is there a way to save a checkpoint in multiple steps so that their PyTrees are combined?

state1 = {'a': jnp.array([1])}
state2 = {'b': jnp.array([2])}

mngr.save(0, args=ocp.args.StandardSave(state1))
mngr.save(0, args=ocp.args.StandardSave(state2))

print(mngr.restore(0))

# I would like to get: { 'a': [1], 'b': [2] }

Motivation: I have a massive checkpoint from Hugging face and I would like to convert that to Orbax. Naive loading and saving the checkpoint will not fit into memory, so I would like to load and save each tensor individually.

[cpgaffney1]

Unfortunately not, though this is a feature we have planned, and it will basically do exactly what you've described.

For now, your best option is to save separate checkpoints and merge them together when restoring.

[jaro-sevcik]

Thank you for the quick reply!

I do not control the loading side (MaxText), so it is a bit unfortunate...